● Harden form-fill analyser with output validators, per-call timeouts, fallback telemetry; raise chunker threshold to 1500

This commit is contained in:
Reece
2026-05-08 11:10:27 +01:00
parent 273865217e
commit c93891dfe0
13 changed files with 1468 additions and 201 deletions
@@ -22,7 +22,9 @@ public enum AiWorkflowOutcome {
CANNOT_CONTINUE("cannot_continue"),
KNOWLEDGE_UPDATE("knowledge_update"),
MULTI_PROFILE_EXTRACTION("multi_profile_extraction"),
BATCH_FILL_RESULT("batch_fill_result");
BATCH_FILL_RESULT("batch_fill_result"),
FORM_ANALYSIS("form_analysis"),
FORM_ANALYSIS_AMBIGUOUS("form_analysis_ambiguous");
private final String value;
+2 -2
View File
@@ -7,12 +7,12 @@
# Configure the model strings passed to pydantic-ai. Provider credentials are handled by
# pydantic-ai and should be set using the provider's native environment variables, for example
# ANTHROPIC_API_KEY or OPENAI_API_KEY.
STIRLING_SMART_MODEL=anthropic:claude-haiku-4-5
STIRLING_SMART_MODEL=anthropic:claude-sonnet-4-6
STIRLING_FAST_MODEL=anthropic:claude-haiku-4-5
# Default output token limits applied by the engine for each model tier.
STIRLING_SMART_MODEL_MAX_TOKENS=8192
STIRLING_FAST_MODEL_MAX_TOKENS=2048
STIRLING_FAST_MODEL_MAX_TOKENS=4096
# RAG Configuration — retrieval-augmented generation is always on.
# Embedding provider credentials are handled natively (e.g. VOYAGE_API_KEY for VoyageAI).
+773 -79
View File
@@ -1,136 +1,830 @@
"""Form Analyser — multi-phase pipeline mirroring PdfEditAgent.
Pipeline:
1. Deterministic prefix pre-pass (no LLM)
- Scans field names for known role prefixes (client_*, contractor_*, etc.)
- Pre-assigns role for any field whose name unambiguously declares one.
- Pattern: ArithmeticScanner in MathAuditorAgent — deterministic-first.
2. Phase 1 — RoleDetectorAgent (smart_model)
- Identifies the *set* of roles per file, cross-file role merges,
cleaned labels, system/skipped fields.
- Does NOT assign individual fields to roles.
- Can refuse via FormAnalysisAmbiguousResponse.
- Pattern: PdfEditAgent.selection_agent — plan only, no params.
3. Phase 2 — FieldAssignerAgent (smart_model)
- For fields not pre-assigned in step 1, decides which role each belongs
to. One batched call per file given the role list from phase 1.
- Pattern: PdfEditParameterSelector — fills in details given the plan.
4. Merge — combines deterministic + LLM assignments into the final
FormAnalysisResponse. Deterministic assignments always win.
"""
from __future__ import annotations
from pydantic_ai import Agent
import asyncio
import logging
import re
from collections.abc import Iterable
from dataclasses import dataclass
from typing import Literal
from pydantic import Field
from pydantic_ai import Agent, ModelRetry, RunContext
from pydantic_ai.output import NativeOutput
from stirling.contracts import format_conversation_history
from stirling.contracts.form_fill import (
AnalysedFileResult,
CleanedLabel,
CrossFileRole,
DetectedRole,
FileFieldSet,
FormAnalysisAmbiguousResponse,
FormAnalysisRequest,
FormAnalysisResponse,
FormAnalysisWorkflowResponse,
FormField,
)
from stirling.models import ApiModel
from stirling.services import AppRuntime
FORM_ANALYSER_SYSTEM_PROMPT = """\
You are a PDF form analyser. You receive field metadata from one or more PDF forms.
logger = logging.getLogger(__name__)
Your job (analysis ONLY — do NOT fill any values):
# ---------------------------------------------------------------------------
# Step 1: Deterministic prefix pre-pass
# ---------------------------------------------------------------------------
1. PER-FILE ROLE DETECTION: For each file, identify distinct sections and the role each \
represents. Use field name prefixes (e.g. "Client" vs "Beneficiary"), nearby page text \
section headers, and form context to group fields by role. Mark which role is the \
"primary person" (the form submitter). If a form has only one section or no section \
distinction, create a single role for it and mark it as primary.
# Role keywords that, when present as a prefix or early token in a field name,
# unambiguously declare which role the field belongs to. Matched against
# tokenised field names — see _deterministic_role_assignments.
ROLE_KEYWORDS: dict[str, str] = {
"client": "Client",
"contractor": "Contractor",
"applicant": "Applicant",
"buyer": "Buyer",
"seller": "Seller",
"employer": "Employer",
"employee": "Employee",
"landlord": "Landlord",
"tenant": "Tenant",
"borrower": "Borrower",
"lender": "Lender",
"guarantor": "Guarantor",
"beneficiary": "Beneficiary",
"policyholder": "Policyholder",
"patient": "Patient",
"spouse": "Spouse",
"witness": "Witness",
"supplier": "Supplier",
"vendor": "Vendor",
"customer": "Customer",
}
PER-FIELD ASSIGNMENT (this is the important bit — do it carefully, don't shortcut):
For every fillable field, decide which role it most plausibly belongs to. Use, in priority order:
_TOKEN_SPLIT = re.compile(r"[_\-\.\[\]\(\)\s/]+")
a) Field-name signal — prefixes/substrings like "client_*", "contractor_*", "applicant_*", \
"buyer_signature", "f_employer_3" carry the role directly.
b) Field-label signal — "Client signature" → Client; "Contractor printed name" → Contractor.
c) Page-text proximity — a "Date" field appearing immediately under a "Client signature" \
line belongs to the Client; the same label under "Contractor signature" belongs to the \
Contractor.
d) Section headers — fields inside a "Client Information" block belong to Client even if \
their own name/label is generic (e.g. "Phone").
For form-level fields that genuinely apply to both/neither party (project description, fees, \
deliverables, invoice frequency, payment terms, contract dates that aren't a signature date): \
assign to the role most directly responsible for that information. Examples — \
"Description of services" → the service provider (Contractor). "Fee amount" → Contractor. \
"Client account number" → Client. "Site address" on a job-site form → the site/project role. \
Use semantic judgement; do not just dump everything on the primary role.
def _deterministic_role_assignments(fields: Iterable[FormField]) -> dict[str, str]:
"""Return {field_name: role_label} for fields whose name unambiguously
declares a role via a known prefix / early token.
ONLY when (a)–(d) and semantic judgement leave a field genuinely ambiguous AND it's a generic \
form-wide field (e.g. an unmarked "Date" with no nearby signature line, "Reference number"): \
fall back to the primary role.
A field is assigned only when exactly one role keyword appears in its
first three tokens. Fields whose name contains zero or multiple role
keywords are left for the LLM to decide.
"""
assignments: dict[str, str] = {}
for field in fields:
if field.read_only:
continue
tokens = [t.lower() for t in _TOKEN_SPLIT.split(field.name) if t]
if not tokens:
continue
head = tokens[:3]
matches = {ROLE_KEYWORDS[t] for t in head if t in ROLE_KEYWORDS}
if len(matches) == 1:
assignments[field.name] = next(iter(matches))
return assignments
EXHAUSTIVENESS: every fillable field must end up in either a role's field list or \
skipped_field_names. No orphans. The downstream filler can't see fields you don't list.
2. CROSS-FILE ROLE MERGING: When the same conceptual role appears in multiple files \
(e.g. "Client" in file A and "Applicant" in file B both mean the primary form submitter), \
merge them into a single cross_file_role with a canonical label. Use semantic understanding \
— "Client", "Applicant", "Account Holder" may all be the same role. Only merge when \
confident they refer to the same type of person/entity.
# ---------------------------------------------------------------------------
# Phase 1: Role detection (no per-field assignment)
# ---------------------------------------------------------------------------
3. LABEL CLEANUP: For fields with unreadable labels (numeric codes like "0021-6009", \
technical IDs like "fld_x7q", names with brackets or underscores), find the real label \
from the nearby page text. Use page text verbatim. Do not clean labels that are already \
clear readable text.
4. INTERNAL FIELD DETECTION: Identify system fields (form IDs, submission references, \
tracking codes, barcodes, buttons, submit buttons). Return their field names in \
skipped_field_names per file.
class _PerFileRoleSummary(ApiModel):
file_id: str
role_labels: list[str]
primary_role_label: str
cleaned_labels: list[CleanedLabel] = Field(default_factory=list)
skipped_field_names: list[str] = Field(default_factory=list)
5. CONFIDENCE SCORING: For each cross_file_role where is_primary_person is true, the caller \
uses the role structure to determine confidence. Set is_primary_person accurately based on:
- Clear primary-person keywords (Client, Applicant, Employee, Patient, Account Holder, \
Policyholder, Tenant, Borrower) → mark as primary.
- If ambiguous (Party A/B, Person 1/2), mark the first/most prominent section as primary.
- If only one section exists, mark it as primary.
STRICT RULES:
- Do NOT return any fill values. This is analysis only.
- Every file must have at least one detected role.
class _CrossFileRoleSummary(ApiModel):
canonical_label: str
file_role_labels: dict[str, str] # file_id -> per-file role label
is_primary_person: bool
class _RoleDetectionPlan(ApiModel):
"""Phase-1 output. Lists roles but does NOT assign fields to them.
The `outcome` literal mirrors the discriminator pattern used by
PdfEditPlanSelection in `pdf_edit.py` — every member of a NativeOutput
union needs a tag so pydantic-ai can disambiguate without ambiguity.
"""
outcome: Literal["role_plan"] = "role_plan"
per_file: list[_PerFileRoleSummary]
cross_file_roles: list[_CrossFileRoleSummary]
message: str = ""
ROLE_DETECTOR_SYSTEM_PROMPT = """\
You analyse PDF forms to identify roles (sections like "Client", "Contractor", \
"Applicant"). Do NOT assign individual fields to roles in this stage — that's a \
separate agent. Your job:
1. PER-FILE ROLES: Identify each distinct role/section in each file. Use field \
name prefixes ("client_*", "f_employer_3"), nearby page-text section headers, and \
form context. If a form has only one section, return one role for it.
2. PRIMARY-PERSON: Mark which role is the form submitter. Heuristics: \
"Client", "Applicant", "Employee", "Patient", "Account Holder", "Policyholder", \
"Tenant", "Borrower", "Buyer" → primary. If ambiguous, mark the first/most \
prominent section. If only one section, mark it primary.
3. CROSS-FILE MERGING: When the same conceptual role appears in multiple files \
(e.g. "Client" in file A and "Applicant" in file B both mean the form submitter), \
merge them into a single cross_file_role with a canonical label. Only merge when \
confident.
4. LABEL CLEANUP: For fields with unreadable labels (numeric codes like \
"0021-6009", technical IDs like "fld_x7q", names with brackets/underscores), find \
the real label from nearby page text and return it as a cleaned_label. Use page \
text verbatim. Skip labels that are already readable.
5. SYSTEM-FIELD DETECTION: Identify form IDs, submission references, tracking \
codes, barcodes, submit buttons. Return their field names in skipped_field_names.
REFUSAL: If the form genuinely has no detectable structure, you may indicate \
ambiguity by returning a role plan with a clear `message` explaining the \
problem (e.g. "47 disconnected fields with no section headers"). The downstream \
agent will surface this to the user.
STRICT:
- Every file must have at least one role.
- cross_file_roles must cover ALL detected roles across all files.
- If a role only appears in one file, it still appears in cross_file_roles with one file_id.
- field_names_by_file MUST be populated — for every detected role, the dict must contain \
each file_id that role appears in, mapping to the list of field names belonging to that \
role in that file. Empty dicts and missing keys break the downstream filler.
- Every fillable field name from the prompt must appear in either some role's field list \
or skipped_field_names. No orphaned fields.
- A role appearing in only one file still appears in cross_file_roles.
- Do NOT return field assignments. That's the next agent's job.
"""
# ---------------------------------------------------------------------------
# Phase 2: Field assignment (given the role plan from phase 1)
# ---------------------------------------------------------------------------
class _FieldAssignment(ApiModel):
field_name: str
role_label: str # must match one of the roles from phase 1
class _FieldAssignmentBatch(ApiModel):
assignments: list[_FieldAssignment]
FIELD_ASSIGNER_SYSTEM_PROMPT = """\
You receive (a) a list of roles already identified in a PDF form and (b) a list \
of fields that need to be assigned to those roles. Your job is to decide, for \
each field, which role it belongs to.
Use, in priority order:
a) Field-name signal — prefixes/substrings carry the role directly.
b) Field-label signal — "Client signature" → Client.
c) Page-text proximity — a "Date" under "Client signature" → Client.
d) Section headers — fields inside a "Client Information" block belong to \
Client even if their own label is generic.
For form-level fields that genuinely apply to both/neither party (project \
description, fees, deliverables, invoice frequency, payment terms, contract \
dates that aren't a signature date): assign to the role most directly responsible \
for that information. Examples — "Description of services" → service provider \
(Contractor); "Fee amount" → Contractor; "Client account number" → Client; \
"Site address" on a job-site form → site/project role. Use semantic judgement; \
don't dump everything on the primary role.
ONLY when (a)–(d) and semantic judgement leave a field genuinely ambiguous AND \
it's a generic form-wide field (e.g. an unmarked "Date" with no nearby signature \
line, "Reference number"): fall back to the primary role.
STRICT:
- Every field given to you must appear in the output exactly once.
- role_label must be one of the roles listed in the prompt.
- Do not invent new roles.
"""
# ---------------------------------------------------------------------------
# Constants used by both phases
# ---------------------------------------------------------------------------
MAX_PAGE_TEXT_CHARS = 1500
# Per-call wall-clock ceilings. Stuck calls would otherwise hold a phase-2
# semaphore slot indefinitely. Retries from output_validator stay inside the
# same Agent.run, so these bound total time including retries.
_PHASE1_TIMEOUT_SECONDS = 120
_PHASE2_TIMEOUT_SECONDS = 60
# ---------------------------------------------------------------------------
# Output validator deps
# ---------------------------------------------------------------------------
@dataclass
class _RoleDetectorDeps:
file_ids: set[str]
@dataclass
class _FieldAssignerDeps:
file_id: str
expected_field_names: set[str]
allowed_role_labels: set[str]
def _validate_role_plan_output(
output: _RoleDetectionPlan | FormAnalysisAmbiguousResponse,
deps: _RoleDetectorDeps,
) -> _RoleDetectionPlan | FormAnalysisAmbiguousResponse:
"""Pure validator for the phase-1 output. Raises ModelRetry on invariant
violations so pydantic-ai prompts the LLM to correct itself instead of
leaving the assemble step to patch around the issue. Extracted as a
module-level function for unit-testability.
"""
if isinstance(output, FormAnalysisAmbiguousResponse):
return output
plan_files = {pf.file_id for pf in output.per_file}
missing = deps.file_ids - plan_files
if missing:
raise ModelRetry(
"Your role plan is missing per_file entries for these file_ids: "
f"{sorted(missing)}. Every request file must appear in per_file."
)
extra = plan_files - deps.file_ids
if extra:
raise ModelRetry(
"Your role plan includes per_file entries for unknown file_ids: "
f"{sorted(extra)}. Use the exact file_ids from the request."
)
for pf in output.per_file:
if not pf.role_labels:
raise ModelRetry(
f"File {pf.file_id} has empty role_labels. Every file must "
f"have at least one role."
)
if pf.primary_role_label not in pf.role_labels:
raise ModelRetry(
f"File {pf.file_id}'s primary_role_label "
f"'{pf.primary_role_label}' is not in role_labels "
f"{pf.role_labels}. The primary must be one of the listed roles."
)
return output
def _validate_field_assignments_output(
output: _FieldAssignmentBatch,
deps: _FieldAssignerDeps,
) -> _FieldAssignmentBatch:
"""Pure validator for the phase-2 output. Each input field must appear
exactly once with a known role label.
"""
seen: set[str] = set()
for fa in output.assignments:
if fa.field_name in seen:
raise ModelRetry(
f"Field '{fa.field_name}' appears in your assignments more "
f"than once. Each field must appear exactly once."
)
seen.add(fa.field_name)
if fa.field_name not in deps.expected_field_names:
raise ModelRetry(
f"Your assignments include unknown field '{fa.field_name}'. "
f"Only assign fields from the list given to you."
)
if fa.role_label not in deps.allowed_role_labels:
raise ModelRetry(
f"Field '{fa.field_name}' was assigned to role "
f"'{fa.role_label}', which is not in the allowed roles "
f"{sorted(deps.allowed_role_labels)}. Use only the roles listed."
)
missing = deps.expected_field_names - seen
if missing:
raise ModelRetry(
f"Your assignments are missing these fields: {sorted(missing)}. "
f"Every field given to you must appear in the output exactly once."
)
return output
# ---------------------------------------------------------------------------
# Orchestrator
# ---------------------------------------------------------------------------
class FormAnalyserAgent:
# Bound concurrent phase-2 LLM calls so a 50-file batch doesn't saturate
# the provider — same pattern as MathAuditorAgent._llm_semaphore.
_PHASE2_CONCURRENCY = 6
def __init__(self, runtime: AppRuntime) -> None:
self.runtime = runtime
self.agent = Agent(
self.role_detector = Agent(
model=runtime.smart_model,
output_type=NativeOutput(FormAnalysisResponse),
system_prompt=FORM_ANALYSER_SYSTEM_PROMPT,
output_type=NativeOutput([_RoleDetectionPlan, FormAnalysisAmbiguousResponse]),
system_prompt=ROLE_DETECTOR_SYSTEM_PROMPT,
model_settings=runtime.smart_model_settings,
deps_type=_RoleDetectorDeps,
)
self.field_assigner = Agent(
model=runtime.smart_model,
output_type=NativeOutput(_FieldAssignmentBatch),
system_prompt=FIELD_ASSIGNER_SYSTEM_PROMPT,
model_settings=runtime.smart_model_settings,
deps_type=_FieldAssignerDeps,
)
self._register_output_validators()
def _register_output_validators(self) -> None:
"""Wire the module-level validator functions onto the two agents.
The LLM gets a chance to correct itself before the assemble step has
to patch around an invariant violation.
"""
@self.role_detector.output_validator
async def _role_plan_validator(
ctx: RunContext[_RoleDetectorDeps],
output: _RoleDetectionPlan | FormAnalysisAmbiguousResponse,
) -> _RoleDetectionPlan | FormAnalysisAmbiguousResponse:
return _validate_role_plan_output(output, ctx.deps)
@self.field_assigner.output_validator
async def _field_assignments_validator(
ctx: RunContext[_FieldAssignerDeps],
output: _FieldAssignmentBatch,
) -> _FieldAssignmentBatch:
return _validate_field_assignments_output(output, ctx.deps)
async def analyse(self, request: FormAnalysisRequest) -> FormAnalysisWorkflowResponse:
# Files with zero form fields can't be analysed — short-circuit upstream
# so phase 1 isn't asked to invent roles for an empty list.
request = request.model_copy(
update={"files": [f for f in request.files if f.form_fields]}
)
if not request.files:
return FormAnalysisAmbiguousResponse(
reason="None of the provided files contain any fillable form fields.",
suggestion="Check that the PDFs are AcroForm documents (not scans).",
)
# Phase 0 — deterministic prefix pre-pass per file.
deterministic: dict[str, dict[str, str]] = {
f.file_id: _deterministic_role_assignments(f.form_fields) for f in request.files
}
# Phase 1 — role detection only.
plan_or_refusal = await self._detect_roles(request, deterministic)
if isinstance(plan_or_refusal, FormAnalysisAmbiguousResponse):
return plan_or_refusal
plan = plan_or_refusal
# Phase 2 — assign remaining (non-deterministic) fields. Files are
# independent so fan them out under a semaphore (mirrors
# ledger/agent.py's asyncio.gather + Semaphore pattern).
per_file_assignments, phase2_warnings = await self._run_phase_2(
request, plan, deterministic
)
async def analyse(self, request: FormAnalysisRequest) -> FormAnalysisResponse:
prompt = self._build_prompt(request)
result = await self.agent.run(prompt)
# Phase 3 — assemble the response in the legacy shape.
return self._assemble_response(
request, plan, per_file_assignments, deterministic, phase2_warnings
)
# ------------------------------------------------------------------
# Phase 2 orchestration
# ------------------------------------------------------------------
async def _run_phase_2(
self,
request: FormAnalysisRequest,
plan: _RoleDetectionPlan,
deterministic: dict[str, dict[str, str]],
) -> tuple[dict[str, dict[str, str]], list[str]]:
"""Returns (per_file_assignments, warnings). Warnings describe per-file
failures so the user can see when phase-2 LLM calls degraded."""
sem = asyncio.Semaphore(self._PHASE2_CONCURRENCY)
async def assign_one(
file_set: FileFieldSet,
file_plan: _PerFileRoleSummary,
) -> tuple[str, dict[str, str]]:
file_det = deterministic.get(file_set.file_id, {})
skipped = set(file_plan.skipped_field_names)
unassigned = [
f for f in file_set.form_fields
if f.name not in file_det and f.name not in skipped and not f.read_only
]
if not unassigned:
return file_set.file_id, dict(file_det)
async with sem:
assigned = await self._assign_fields(file_set, file_plan, unassigned)
merged = dict(file_det)
for fa in assigned.assignments:
# Deterministic always wins.
if fa.field_name not in merged and fa.role_label in file_plan.role_labels:
merged[fa.field_name] = fa.role_label
return file_set.file_id, merged
tasks: list[tuple[str, str]] = [] # (file_id, file_name) for warnings
coros = []
for file_set in request.files:
file_plan = next((pf for pf in plan.per_file if pf.file_id == file_set.file_id), None)
if file_plan is None:
continue
tasks.append((file_set.file_id, file_set.file_name))
coros.append(assign_one(file_set, file_plan))
results = await asyncio.gather(*coros, return_exceptions=True)
per_file_assignments: dict[str, dict[str, str]] = {}
warnings: list[str] = []
for (file_id, file_name), outcome in zip(tasks, results, strict=True):
if isinstance(outcome, BaseException):
# Phase-2 failures degrade gracefully — the assemble step's
# fallback ("everything else → primary") still produces a
# usable response. Surface a warning so the user knows.
logger.warning(
"form_analyser.fallback",
extra={
"layer": "phase2_call_failed",
"file_id": file_id,
"error_type": type(outcome).__name__,
},
)
warnings.append(
f"{file_name}: field-assignment AI call failed ({type(outcome).__name__}); "
f"fields attached to primary role as a fallback."
)
continue
file_id_out, merged = outcome
per_file_assignments[file_id_out] = merged
return per_file_assignments, warnings
# ------------------------------------------------------------------
# Phase 1 helpers
# ------------------------------------------------------------------
async def _detect_roles(
self,
request: FormAnalysisRequest,
deterministic: dict[str, dict[str, str]],
) -> _RoleDetectionPlan | FormAnalysisAmbiguousResponse:
prompt = self._build_detection_prompt(request, deterministic)
deps = _RoleDetectorDeps(file_ids={f.file_id for f in request.files})
result = await asyncio.wait_for(
self.role_detector.run(prompt, deps=deps),
timeout=_PHASE1_TIMEOUT_SECONDS,
)
return result.output
def _build_prompt(self, request: FormAnalysisRequest) -> str:
def _build_detection_prompt(
self,
request: FormAnalysisRequest,
deterministic: dict[str, dict[str, str]],
) -> str:
sections = []
for file_set in request.files:
page_texts: dict[str, int] = {}
field_page_map: dict[str, int] = {}
for field in file_set.form_fields:
if field.nearby_page_text:
text = field.nearby_page_text
if text not in page_texts:
page_texts[text] = len(page_texts)
field_page_map[field.name] = page_texts[text]
page_section = "\n".join(
f" [Page {pid}]: {text[:MAX_PAGE_TEXT_CHARS]}" for text, pid in page_texts.items()
page_section, field_page_map = self._format_page_texts(file_set)
fields_text = "\n".join(
self._format_field(f, field_page_map.get(f.name)) for f in file_set.form_fields
)
fields_text = "\n".join(self._format_field(f, field_page_map.get(f.name)) for f in file_set.form_fields)
det_section = self._format_deterministic_hints(deterministic.get(file_set.file_id, {}))
sections.append(
f"=== FILE: {file_set.file_name} (id={file_set.file_id}) ===\n"
f"Page texts:\n{page_section}\n\n"
f"{det_section}"
f"Fields:\n{fields_text}"
)
history = format_conversation_history(request.conversation_history)
return (
f"Conversation history:\n{history}\n\n"
f"Analyse {len(request.files)} form(s):\n\n"
f"Identify roles in {len(request.files)} form(s). "
"Do NOT assign individual fields to roles — that's a separate agent.\n\n"
+ "\n\n".join(sections)
+ "\n\nDetect roles per file, merge matching roles across files, "
"clean unreadable labels, identify internal/system fields."
)
def _format_deterministic_hints(self, assignments: dict[str, str]) -> str:
"""Build a "Pre-detected roles" prompt block listing the regex-derived
role assignments. Importantly, lists each role's fields so the LLM has
evidence rather than just a label, and asks it to use these labels
verbatim — otherwise role names diverge between phases and the
downstream merge silently drops fields.
"""
if not assignments:
return ""
by_role: dict[str, list[str]] = {}
for field_name, role_label in assignments.items():
by_role.setdefault(role_label, []).append(field_name)
lines = [
"Pre-detected roles in this file (from field-name prefixes — keep these "
"exact role labels in your `role_labels` output; do not rename them):",
]
for role_label in sorted(by_role):
fields = ", ".join(sorted(by_role[role_label])[:8])
extra = "" if len(by_role[role_label]) <= 8 else f" (+{len(by_role[role_label]) - 8} more)"
lines.append(f" - {role_label}: {fields}{extra}")
return "\n".join(lines) + "\n\n"
# ------------------------------------------------------------------
# Phase 2 helpers
# ------------------------------------------------------------------
async def _assign_fields(
self,
file_set: FileFieldSet,
file_plan: _PerFileRoleSummary,
unassigned: list[FormField],
) -> _FieldAssignmentBatch:
prompt = self._build_assignment_prompt(file_set, file_plan, unassigned)
deps = _FieldAssignerDeps(
file_id=file_set.file_id,
expected_field_names={f.name for f in unassigned},
allowed_role_labels=set(file_plan.role_labels),
)
result = await asyncio.wait_for(
self.field_assigner.run(prompt, deps=deps),
timeout=_PHASE2_TIMEOUT_SECONDS,
)
return result.output
def _build_assignment_prompt(
self,
file_set: FileFieldSet,
file_plan: _PerFileRoleSummary,
unassigned: list[FormField],
) -> str:
page_section, field_page_map = self._format_page_texts(file_set)
fields_text = "\n".join(
self._format_field(f, field_page_map.get(f.name)) for f in unassigned
)
roles_list = ", ".join(file_plan.role_labels)
return (
f"=== FILE: {file_set.file_name} (id={file_set.file_id}) ===\n"
f"Roles in this file: {roles_list}\n"
f"Primary role: {file_plan.primary_role_label}\n\n"
f"Page texts:\n{page_section}\n\n"
f"Assign each of these {len(unassigned)} field(s) to one of the roles above:\n"
f"{fields_text}"
)
# ------------------------------------------------------------------
# Phase 3 — assemble the legacy-shaped FormAnalysisResponse
# ------------------------------------------------------------------
def _assemble_response(
self,
request: FormAnalysisRequest,
plan: _RoleDetectionPlan,
per_file_assignments: dict[str, dict[str, str]],
deterministic: dict[str, dict[str, str]],
extra_warnings: list[str] | None = None,
) -> FormAnalysisResponse:
per_file_results: list[AnalysedFileResult] = []
warnings: list[str] = list(extra_warnings or [])
for file_set in request.files:
file_plan = next((pf for pf in plan.per_file if pf.file_id == file_set.file_id), None)
if file_plan is None:
logger.warning(
"form_analyser.fallback",
extra={
"layer": "phase1_missing_per_file",
"file_id": file_set.file_id,
},
)
warnings.append(
f"{file_set.file_name}: role detector returned no roles — falling back "
f"to a single Primary role."
)
file_plan = _PerFileRoleSummary(
file_id=file_set.file_id,
role_labels=["Primary"],
primary_role_label="Primary",
)
# Defensive copy — _assemble_response shouldn't mutate dicts owned
# by phase-2's result. (Today the caller doesn't reuse them, but
# callers shouldn't have to remember that.)
assignments = dict(per_file_assignments.get(file_set.file_id, {}))
skipped = set(file_plan.skipped_field_names)
# Anything still missing — the LLM didn't assign and there was no
# deterministic prefix — gets attached to the primary role.
orphan_count = 0
for field in file_set.form_fields:
if field.read_only or field.name in skipped or field.name in assignments:
continue
assignments[field.name] = file_plan.primary_role_label
orphan_count += 1
if orphan_count:
logger.warning(
"form_analyser.fallback",
extra={
"layer": "orphan_to_primary",
"file_id": file_set.file_id,
"count": orphan_count,
},
)
# Union the LLM's role labels with any role label that appears in
# `assignments` (e.g. from the deterministic pre-pass) but wasn't
# listed by phase 1. Without this, deterministic-tagged fields under
# a role the LLM renamed/missed would be silently dropped from every
# DetectedRole. Order: LLM's order first (preserves primary), then
# any extras in the order they first appear in assignments.
extra_labels: list[str] = []
seen = set(file_plan.role_labels)
for label in assignments.values():
if label not in seen:
seen.add(label)
extra_labels.append(label)
all_role_labels = [*file_plan.role_labels, *extra_labels]
if extra_labels:
logger.warning(
"form_analyser.fallback",
extra={
"layer": "deterministic_label_union",
"file_id": file_set.file_id,
"extra_labels": extra_labels,
},
)
warnings.append(
f"{file_set.file_name}: kept {len(extra_labels)} pre-detected "
f"role(s) not listed by the AI: {', '.join(extra_labels)}."
)
detected: list[DetectedRole] = []
for role_label in all_role_labels:
role_fields = [n for n, r in assignments.items() if r == role_label]
detected.append(
DetectedRole(
role_label=role_label,
field_names=role_fields,
is_primary_person=(role_label == file_plan.primary_role_label),
)
)
per_file_results.append(
AnalysedFileResult(
file_id=file_set.file_id,
file_name=file_set.file_name,
detected_roles=detected,
cleaned_labels=file_plan.cleaned_labels,
skipped_field_names=file_plan.skipped_field_names,
)
)
cross_file_roles = self._build_cross_file_roles(plan, per_file_results)
det_count = sum(len(d) for d in deterministic.values())
message_parts: list[str] = []
if plan.message:
message_parts.append(plan.message)
if det_count:
message_parts.append(
f"{det_count} field(s) assigned deterministically by name prefix."
)
message_parts.extend(warnings)
return FormAnalysisResponse(
per_file=per_file_results,
cross_file_roles=cross_file_roles,
message=" ".join(message_parts).strip(),
)
def _build_cross_file_roles(
self,
plan: _RoleDetectionPlan,
per_file_results: list[AnalysedFileResult],
) -> list[CrossFileRole]:
results_by_id = {r.file_id: r for r in per_file_results}
cross: list[CrossFileRole] = []
# Track which (file_id, per_file_role_label) pairs we successfully
# resolved — used both for the safety net and to report which files
# the cross-file role actually covers.
resolved_pairs: list[tuple[str, str]] = []
for cfr in plan.cross_file_roles:
field_names_by_file: dict[str, list[str]] = {}
for file_id, per_file_role_label in cfr.file_role_labels.items():
pf_result = results_by_id.get(file_id)
if pf_result is None:
continue
role = self._resolve_per_file_role(pf_result, per_file_role_label, cfr.canonical_label)
if role is None:
continue
field_names_by_file[file_id] = list(role.field_names)
resolved_pairs.append((file_id, role.role_label))
cross.append(
CrossFileRole(
role_label=cfr.canonical_label,
file_ids=list(field_names_by_file.keys()),
field_names_by_file=field_names_by_file,
is_primary_person=cfr.is_primary_person,
)
)
# Safety net — any per-file role label that didn't make it into a
# cross_file_role at all gets its own single-file entry.
covered: set[tuple[str, str]] = set(resolved_pairs)
for cfr in plan.cross_file_roles:
for fid, lbl in cfr.file_role_labels.items():
covered.add((fid, lbl))
for pf in per_file_results:
for role in pf.detected_roles:
if (pf.file_id, role.role_label) in covered:
continue
cross.append(
CrossFileRole(
role_label=role.role_label,
file_ids=[pf.file_id],
field_names_by_file={pf.file_id: list(role.field_names)},
is_primary_person=role.is_primary_person,
)
)
return cross
def _resolve_per_file_role(
self,
pf_result: AnalysedFileResult,
requested_label: str,
canonical_label: str,
) -> DetectedRole | None:
"""Find the DetectedRole matching `requested_label`, with fallbacks.
The LLM is told to put each file's per-file role label in
`cfr.file_role_labels[file_id]`, but inconsistencies happen — it may
put the canonical label there instead, or use slightly different
capitalisation. This method tries:
1. Exact match on requested_label.
2. Case-insensitive match on requested_label.
3. Exact / case-insensitive match on canonical_label.
Returns None if nothing matches.
"""
# 1. Exact.
for role in pf_result.detected_roles:
if role.role_label == requested_label:
return role
# 2. Case-insensitive on requested_label.
requested_ci = requested_label.casefold()
for role in pf_result.detected_roles:
if role.role_label.casefold() == requested_ci:
return role
# 3. Fall back to canonical label (exact, then case-insensitive).
if canonical_label != requested_label:
for role in pf_result.detected_roles:
if role.role_label == canonical_label:
return role
canonical_ci = canonical_label.casefold()
for role in pf_result.detected_roles:
if role.role_label.casefold() == canonical_ci:
return role
return None
# ------------------------------------------------------------------
# Shared formatting helpers
# ------------------------------------------------------------------
def _format_page_texts(self, file_set: FileFieldSet) -> tuple[str, dict[str, int]]:
page_texts: dict[str, int] = {}
field_page_map: dict[str, int] = {}
for field in file_set.form_fields:
if field.nearby_page_text:
text = field.nearby_page_text
if text not in page_texts:
page_texts[text] = len(page_texts)
field_page_map[field.name] = page_texts[text]
page_section = "\n".join(
f" [Page {pid}]: {text[:MAX_PAGE_TEXT_CHARS]}" for text, pid in page_texts.items()
)
return page_section, field_page_map
def _format_field(self, field: FormField, page_id: int | None) -> str:
parts = [f"- name={field.name}, type={field.type}"]
if field.label:
+27 -1
View File
@@ -22,10 +22,28 @@ STRICT RULES:
- ONLY fill a field if a knowledge entry is clearly the same information.
- For checkbox fields: value must be "Yes" or "Off".
- For radio/combobox/listbox fields: value MUST be one of the field's options.
- Skip read-only fields.
- For multiline fields, you may emit multi-line values (e.g. addresses).
- For multi_select fields, the value should be a comma-separated list of options.
- Do NOT perform role detection — it is already done.
- Return filled_fields per file using the file_id from the request.
- If no fields match for a file, return an empty filled_fields list for that file.
MULTI-ENTITY KNOWLEDGE PREFERENCE:
When the dictionary has BOTH a plain key (e.g. `first_name`) and a role-prefixed \
version (e.g. `client_first_name`), the role-prefixed key is authoritative for \
this file's role_label — prefer it. The plain key is a shared/fallback value used \
when the role-prefixed version is missing.
NEARBY PAGE TEXT:
A field's `nearby_page_text` snippet is the literal text surrounding the field on \
the page. Use it to disambiguate generic labels — e.g. a field labelled "Date" with \
"Contract effective date:" in its nearby text wants `start_date`, not \
`date_of_birth`.
EXISTING VALUE:
A field's `value=` shows what's already filled in. If the existing value already \
matches the right knowledge entry, you may still emit it (the caller decides \
whether to apply). Do NOT overwrite a non-empty existing value with a stale guess.
"""
@@ -67,10 +85,18 @@ class FormFillerAgent:
parts = [f"- name={field.name}, type={field.type}"]
if field.label:
parts.append(f" label={field.label}")
if field.value:
parts.append(f" value={field.value}")
if field.tooltip:
parts.append(f" tooltip={field.tooltip}")
if field.options:
parts.append(f" options={field.options}")
if field.required:
parts.append(" required=true")
if field.multiline:
parts.append(" multiline=true")
if field.multi_select:
parts.append(" multi_select=true")
if field.nearby_page_text:
parts.append(f" nearby_page_text={field.nearby_page_text}")
return "\n".join(parts)
+3 -3
View File
@@ -14,7 +14,7 @@ from stirling.contracts.form_fill import (
DocumentExtractionRequest,
DocumentExtractionResponse,
FormAnalysisRequest,
FormAnalysisResponse,
FormAnalysisWorkflowResponse,
FormFillBatchRequest,
FormFillBatchResponse,
)
@@ -22,11 +22,11 @@ from stirling.contracts.form_fill import (
router = APIRouter(prefix="/api/v1/form/ai", tags=["form-fill"])
@router.post("/analyse", response_model=FormAnalysisResponse)
@router.post("/analyse", response_model=FormAnalysisWorkflowResponse)
async def analyse_forms(
request: FormAnalysisRequest,
agent: Annotated[FormAnalyserAgent, Depends(get_form_analyser_agent)],
) -> FormAnalysisResponse:
) -> FormAnalysisWorkflowResponse:
return await agent.analyse(request)
+2
View File
@@ -61,6 +61,8 @@ class WorkflowOutcome(StrEnum):
KNOWLEDGE_UPDATE = "knowledge_update"
MULTI_PROFILE_EXTRACTION = "multi_profile_extraction"
BATCH_FILL_RESULT = "batch_fill_result"
FORM_ANALYSIS = "form_analysis"
FORM_ANALYSIS_AMBIGUOUS = "form_analysis_ambiguous"
class ArtifactKind(StrEnum):
@@ -119,11 +119,31 @@ class CrossFileRole(ApiModel):
class FormAnalysisResponse(ApiModel):
outcome: Literal[WorkflowOutcome.FORM_ANALYSIS] = WorkflowOutcome.FORM_ANALYSIS
per_file: list[AnalysedFileResult]
cross_file_roles: list[CrossFileRole]
message: str
class FormAnalysisAmbiguousResponse(ApiModel):
"""Returned when the analyser cannot confidently identify roles or fields.
Mirrors the EditClarificationRequest pattern from PdfEditAgent — gives the
LLM an escape valve so it can refuse rather than fabricate role assignments
on inscrutable forms.
"""
outcome: Literal[WorkflowOutcome.FORM_ANALYSIS_AMBIGUOUS] = WorkflowOutcome.FORM_ANALYSIS_AMBIGUOUS
reason: str = Field(max_length=2000)
suggestion: str | None = Field(default=None, max_length=2000)
FormAnalysisWorkflowResponse = Annotated[
FormAnalysisResponse | FormAnalysisAmbiguousResponse,
Field(discriminator="outcome"),
]
# --- Batch Fill ---
+404
View File
@@ -260,3 +260,407 @@ async def test_document_extractor_multi_profile_outcome() -> None:
assert result.outcome == "multi_profile_extraction"
assert len(result.proposed_profiles) == 2
# ---------------------------------------------------------------------------
# Pipeline-piece tests — these exercise the deterministic pre-pass and the
# response assembler directly, without going through the LLM. Stubbing the
# full agent is not enough since the new pipeline has correctness logic
# inside _deterministic_role_assignments and _assemble_response.
# ---------------------------------------------------------------------------
def _make_field(name: str, **kwargs: object) -> FormField:
"""Helper — build a minimal FormField with sensible defaults."""
base: dict[str, object] = {"name": name, "type": "text"}
base.update(kwargs)
return FormField(**base) # type: ignore[arg-type]
def test_deterministic_assigns_obvious_prefix() -> None:
from stirling.agents.form_analyser import _deterministic_role_assignments
fields = [
_make_field("client_first_name"),
_make_field("contractor_signature"),
_make_field("date"), # ambiguous — no role keyword
]
result = _deterministic_role_assignments(fields)
assert result == {"client_first_name": "Client", "contractor_signature": "Contractor"}
def test_deterministic_skips_multi_role_names() -> None:
from stirling.agents.form_analyser import _deterministic_role_assignments
# A name containing two role keywords is ambiguous — leave it for the LLM.
fields = [_make_field("client_contractor_summary")]
assert _deterministic_role_assignments(fields) == {}
def test_deterministic_skips_readonly() -> None:
from stirling.agents.form_analyser import _deterministic_role_assignments
fields = [_make_field("client_first_name", read_only=True)]
assert _deterministic_role_assignments(fields) == {}
def test_deterministic_handles_irs_naming() -> None:
from stirling.agents.form_analyser import _deterministic_role_assignments
# IRS-style names like "f_employer_3" should still resolve to Employer.
fields = [_make_field("f_employer_3"), _make_field("f_employee_1")]
result = _deterministic_role_assignments(fields)
assert result == {"f_employer_3": "Employer", "f_employee_1": "Employee"}
def test_assemble_preserves_deterministic_labels_when_llm_renames() -> None:
"""Regression test for Bug C/H: when the regex pre-pass tags fields under
"Client"/"Contractor" but the LLM phase 1 returns ["Customer", "Vendor"],
the deterministic-tagged fields must still appear in the final response —
they used to be silently dropped from every DetectedRole.
"""
from stirling.agents.form_analyser import (
FormAnalyserAgent,
_PerFileRoleSummary,
_RoleDetectionPlan,
)
runtime = build_runtime(build_app_settings())
agent = FormAnalyserAgent(runtime)
request = FormAnalysisRequest(
files=[
FileFieldSet(
file_id="f1",
file_name="form.pdf",
form_fields=[
_make_field("client_first_name"), # deterministic → Client
_make_field("contractor_signature"), # deterministic → Contractor
_make_field("project_description"), # LLM-assigned
],
)
]
)
plan = _RoleDetectionPlan(
per_file=[
_PerFileRoleSummary(
file_id="f1",
role_labels=["Customer", "Vendor"], # LLM renamed
primary_role_label="Customer",
)
],
cross_file_roles=[],
)
deterministic = {
"f1": {"client_first_name": "Client", "contractor_signature": "Contractor"}
}
per_file_assignments = {
"f1": {
"client_first_name": "Client",
"contractor_signature": "Contractor",
"project_description": "Vendor",
}
}
result = agent._assemble_response(request, plan, per_file_assignments, deterministic)
# Every fillable field must end up in *some* role's field_names.
all_assigned = {
n
for r in result.per_file[0].detected_roles
for n in r.field_names
}
assert all_assigned == {"client_first_name", "contractor_signature", "project_description"}
# Both deterministic labels survive even though they're not in plan.role_labels.
role_labels = {r.role_label for r in result.per_file[0].detected_roles}
assert {"Client", "Contractor", "Customer", "Vendor"} <= role_labels
def test_assemble_attaches_orphan_fields_to_primary() -> None:
"""Phase 2 omits a field. The safety net attaches it to the primary role."""
from stirling.agents.form_analyser import (
FormAnalyserAgent,
_PerFileRoleSummary,
_RoleDetectionPlan,
)
runtime = build_runtime(build_app_settings())
agent = FormAnalyserAgent(runtime)
request = FormAnalysisRequest(
files=[
FileFieldSet(
file_id="f1",
file_name="form.pdf",
form_fields=[
_make_field("known_field"),
_make_field("orphan_field"), # neither LLM nor regex assigned it
],
)
]
)
plan = _RoleDetectionPlan(
per_file=[
_PerFileRoleSummary(
file_id="f1",
role_labels=["Primary"],
primary_role_label="Primary",
)
],
cross_file_roles=[],
)
result = agent._assemble_response(
request,
plan,
per_file_assignments={"f1": {"known_field": "Primary"}},
deterministic={"f1": {}},
)
primary = next(r for r in result.per_file[0].detected_roles if r.is_primary_person)
assert "orphan_field" in primary.field_names
def test_resolve_per_file_role_case_insensitive() -> None:
from stirling.agents.form_analyser import FormAnalyserAgent
runtime = build_runtime(build_app_settings())
agent = FormAnalyserAgent(runtime)
pf = AnalysedFileResult(
file_id="f1",
file_name="form.pdf",
detected_roles=[
DetectedRole(role_label="Client", field_names=["a"], is_primary_person=True),
],
)
# Exact match.
assert agent._resolve_per_file_role(pf, "Client", "Client") is not None
# Different case still matches.
assert agent._resolve_per_file_role(pf, "client", "Client") is not None
# Falls back to canonical when requested label isn't there at all.
assert agent._resolve_per_file_role(pf, "Applicant", "Client") is not None
# Truly unknown role → None.
assert agent._resolve_per_file_role(pf, "Witness", "Witness") is None
# ---------------------------------------------------------------------------
# Output-validator tests — exercised as pure functions so we don't need an
# LLM stub. Each test feeds an LLM response that violates an invariant and
# asserts the validator raises ModelRetry, which is what pydantic-ai uses
# to re-prompt the model.
# ---------------------------------------------------------------------------
def test_role_plan_validator_rejects_missing_file_id() -> None:
from pydantic_ai import ModelRetry
from stirling.agents.form_analyser import (
_PerFileRoleSummary,
_RoleDetectionPlan,
_RoleDetectorDeps,
_validate_role_plan_output,
)
deps = _RoleDetectorDeps(file_ids={"f1", "f2"})
plan = _RoleDetectionPlan(
per_file=[
_PerFileRoleSummary(file_id="f1", role_labels=["Client"], primary_role_label="Client"),
],
cross_file_roles=[],
)
with pytest.raises(ModelRetry, match="missing per_file"):
_validate_role_plan_output(plan, deps)
def test_role_plan_validator_rejects_unknown_file_id() -> None:
from pydantic_ai import ModelRetry
from stirling.agents.form_analyser import (
_PerFileRoleSummary,
_RoleDetectionPlan,
_RoleDetectorDeps,
_validate_role_plan_output,
)
deps = _RoleDetectorDeps(file_ids={"f1"})
plan = _RoleDetectionPlan(
per_file=[
_PerFileRoleSummary(file_id="f1", role_labels=["Client"], primary_role_label="Client"),
_PerFileRoleSummary(file_id="ghost", role_labels=["Client"], primary_role_label="Client"),
],
cross_file_roles=[],
)
with pytest.raises(ModelRetry, match="unknown file_ids"):
_validate_role_plan_output(plan, deps)
def test_role_plan_validator_rejects_invalid_primary_label() -> None:
from pydantic_ai import ModelRetry
from stirling.agents.form_analyser import (
_PerFileRoleSummary,
_RoleDetectionPlan,
_RoleDetectorDeps,
_validate_role_plan_output,
)
deps = _RoleDetectorDeps(file_ids={"f1"})
plan = _RoleDetectionPlan(
per_file=[
_PerFileRoleSummary(
file_id="f1",
role_labels=["Client", "Contractor"],
primary_role_label="Vendor", # not in role_labels
),
],
cross_file_roles=[],
)
with pytest.raises(ModelRetry, match="primary_role_label"):
_validate_role_plan_output(plan, deps)
def test_role_plan_validator_passes_ambiguous_response_through() -> None:
"""Refusal responses bypass the validator — there's no plan to check."""
from stirling.agents.form_analyser import (
_RoleDetectorDeps,
_validate_role_plan_output,
)
from stirling.contracts.form_fill import FormAnalysisAmbiguousResponse
deps = _RoleDetectorDeps(file_ids={"f1"})
refusal = FormAnalysisAmbiguousResponse(reason="cannot tell roles", suggestion=None)
assert _validate_role_plan_output(refusal, deps) is refusal
def test_field_assigner_validator_rejects_unknown_field() -> None:
from pydantic_ai import ModelRetry
from stirling.agents.form_analyser import (
_FieldAssignerDeps,
_FieldAssignment,
_FieldAssignmentBatch,
_validate_field_assignments_output,
)
deps = _FieldAssignerDeps(
file_id="f1",
expected_field_names={"a", "b"},
allowed_role_labels={"Client"},
)
batch = _FieldAssignmentBatch(
assignments=[
_FieldAssignment(field_name="a", role_label="Client"),
_FieldAssignment(field_name="b", role_label="Client"),
_FieldAssignment(field_name="c", role_label="Client"), # not in expected
]
)
with pytest.raises(ModelRetry, match="unknown field 'c'"):
_validate_field_assignments_output(batch, deps)
def test_field_assigner_validator_rejects_unknown_role() -> None:
from pydantic_ai import ModelRetry
from stirling.agents.form_analyser import (
_FieldAssignerDeps,
_FieldAssignment,
_FieldAssignmentBatch,
_validate_field_assignments_output,
)
deps = _FieldAssignerDeps(
file_id="f1",
expected_field_names={"a"},
allowed_role_labels={"Client", "Contractor"},
)
batch = _FieldAssignmentBatch(
assignments=[_FieldAssignment(field_name="a", role_label="Vendor")]
)
with pytest.raises(ModelRetry, match="not in the allowed roles"):
_validate_field_assignments_output(batch, deps)
def test_field_assigner_validator_rejects_missing_field() -> None:
from pydantic_ai import ModelRetry
from stirling.agents.form_analyser import (
_FieldAssignerDeps,
_FieldAssignment,
_FieldAssignmentBatch,
_validate_field_assignments_output,
)
deps = _FieldAssignerDeps(
file_id="f1",
expected_field_names={"a", "b"},
allowed_role_labels={"Client"},
)
batch = _FieldAssignmentBatch(
assignments=[_FieldAssignment(field_name="a", role_label="Client")]
)
with pytest.raises(ModelRetry, match="missing these fields"):
_validate_field_assignments_output(batch, deps)
def test_field_assigner_validator_rejects_duplicate_field() -> None:
from pydantic_ai import ModelRetry
from stirling.agents.form_analyser import (
_FieldAssignerDeps,
_FieldAssignment,
_FieldAssignmentBatch,
_validate_field_assignments_output,
)
deps = _FieldAssignerDeps(
file_id="f1",
expected_field_names={"a"},
allowed_role_labels={"Client"},
)
batch = _FieldAssignmentBatch(
assignments=[
_FieldAssignment(field_name="a", role_label="Client"),
_FieldAssignment(field_name="a", role_label="Client"),
]
)
with pytest.raises(ModelRetry, match="more than once"):
_validate_field_assignments_output(batch, deps)
def test_field_assigner_validator_accepts_valid_output() -> None:
"""Sanity — a well-formed batch returns unchanged."""
from stirling.agents.form_analyser import (
_FieldAssignerDeps,
_FieldAssignment,
_FieldAssignmentBatch,
_validate_field_assignments_output,
)
deps = _FieldAssignerDeps(
file_id="f1",
expected_field_names={"a", "b"},
allowed_role_labels={"Client", "Contractor"},
)
batch = _FieldAssignmentBatch(
assignments=[
_FieldAssignment(field_name="a", role_label="Client"),
_FieldAssignment(field_name="b", role_label="Contractor"),
]
)
assert _validate_field_assignments_output(batch, deps) is batch
@pytest.mark.anyio
async def test_analyser_refuses_when_all_files_empty() -> None:
"""A request with files but zero form fields should surface as ambiguity."""
from stirling.agents.form_analyser import FormAnalyserAgent
from stirling.contracts.form_fill import FormAnalysisAmbiguousResponse
runtime = build_runtime(build_app_settings())
agent = FormAnalyserAgent(runtime)
request = FormAnalysisRequest(
files=[FileFieldSet(file_id="f1", file_name="empty.pdf", form_fields=[])]
)
result = await agent.analyse(request)
assert isinstance(result, FormAnalysisAmbiguousResponse)
@@ -7,7 +7,7 @@ import type {
DocumentExtractionRequest,
DocumentExtractionResponse,
FormAnalysisRequest,
FormAnalysisResponse,
FormAnalysisWorkflowResponse,
FormFillBatchRequest,
FormFillBatchResponse,
} from './types';
@@ -29,8 +29,8 @@ export async function extractFromDocuments(
export async function analyseMultipleForms(
request: FormAnalysisRequest
): Promise<FormAnalysisResponse> {
const response = await aiEngineClient.post<FormAnalysisResponse>(
): Promise<FormAnalysisWorkflowResponse> {
const response = await aiEngineClient.post<FormAnalysisWorkflowResponse>(
'/api/v1/form/ai/analyse',
request
);
@@ -0,0 +1,36 @@
/**
* Shared error helpers for AI Form Fill — both the analyse hook and the
* batch-fill hook used to define the same `describeError` inline.
*/
/**
* Pull the most useful detail out of an axios/fetch error. FastAPI 422s come
* back as `{detail: [{loc, msg, type}, ...]}`; plain HTTP errors carry
* `{detail: string}` or `{message: string}`. Falls back to `e.message` and
* finally to the supplied `fallback` string.
*/
export function describeError(err: unknown, fallback = 'Request failed.'): string {
if (err && typeof err === 'object') {
const e = err as {
response?: { data?: unknown };
message?: unknown;
};
const data = e.response?.data;
if (data && typeof data === 'object') {
const d = data as { detail?: unknown; message?: unknown };
if (Array.isArray(d.detail)) {
return d.detail
.map((entry) => {
const item = entry as { loc?: unknown; msg?: unknown; type?: unknown };
const loc = Array.isArray(item.loc) ? item.loc.join('.') : '';
return `${loc}: ${item.msg ?? item.type ?? ''}`;
})
.join('; ');
}
if (typeof d.detail === 'string') return d.detail;
if (typeof d.message === 'string') return d.message;
}
if (typeof e.message === 'string') return e.message;
}
return fallback;
}
@@ -89,11 +89,22 @@ export interface CrossFileRole {
}
export interface FormAnalysisResponse {
outcome: 'form_analysis';
perFile: AnalysedFileResult[];
crossFileRoles: CrossFileRole[];
message: string;
}
export interface FormAnalysisAmbiguousResponse {
outcome: 'form_analysis_ambiguous';
reason: string;
suggestion: string | null;
}
export type FormAnalysisWorkflowResponse =
| FormAnalysisResponse
| FormAnalysisAmbiguousResponse;
// --- Batch Fill ---
export interface FileFillRequest {
@@ -13,9 +13,12 @@ import { useState, useCallback } from 'react';
import { fillFormFields } from '@app/tools/formFill/formApi';
import { useFileManagement } from '@app/contexts/FileContext';
import { fillFormsBatch } from './aiFormFillApi';
import { describeError as describeErrorShared } from './errorUtils';
import { mergeEntitiesForFill, type MergeResult } from './entityTypes';
import type { Entity } from './entityTypes';
import { resolveDynamicValues } from './workflowTemplates';
const describeError = (err: unknown) => describeErrorShared(err, 'Batch fill failed.');
import type {
FormFillBatchResponse,
FormField,
@@ -111,36 +114,55 @@ const INITIAL_STATE: BatchFillState = {
/** Sanity cap so a misclick doesn't try to fill 10,000 PDFs. */
const MAX_VARIANTS_PER_FILE = 100;
/**
* Mirrors MathAuditorAgent._llm_semaphore (engine/.../ledger/agent.py) — bounds
* how many filler calls can be in flight at once so a 50-variant batch doesn't
* hammer the engine with 50 simultaneous requests.
*/
const MAX_CONCURRENT_FILLER_CALLS = 6;
type Settled<R> = { ok: true; value: R } | { ok: false; error: unknown };
/**
* Bounded-concurrency map that collects per-item outcomes (allSettled-style).
* One failing worker does NOT abort the rest — the caller decides what to do
* with the failures. This matters for batch fills: a single 5xx on variant 7
* shouldn't throw away the 49 successful variants.
*/
async function runWithConcurrency<T, R>(
items: T[],
limit: number,
worker: (item: T, index: number) => Promise<R>,
): Promise<Settled<R>[]> {
const results = new Array<Settled<R>>(items.length);
let next = 0;
const runners = Array.from({ length: Math.min(limit, items.length) }, async () => {
while (true) {
const i = next++;
if (i >= items.length) return;
try {
results[i] = { ok: true, value: await worker(items[i], i) };
} catch (error) {
results[i] = { ok: false, error };
}
}
});
await Promise.all(runners);
return results;
}
// Match the engine contract caps (engine/src/stirling/contracts/form_fill.py).
const MAX_LABEL = 500;
const MAX_TOOLTIP = 1000;
const MAX_NAME = 200;
const MAX_VALUE = 2000;
const MAX_NEARBY_PAGE_TEXT = 8000;
function clamp(s: string | undefined | null, max: number): string | undefined {
if (s == null) return undefined;
return s.length > max ? s.slice(0, max) : s;
}
/** Pull the most useful detail out of an axios/fetch error. */
function describeError(err: unknown): string {
if (err && typeof err === 'object') {
const e = err as any;
const data = e.response?.data;
if (data) {
// FastAPI 422 returns {detail: [{loc, msg, type}, ...]}
if (Array.isArray(data.detail)) {
return data.detail
.map((d: any) => `${(d.loc ?? []).join('.')}: ${d.msg ?? d.type ?? ''}`)
.join('; ');
}
if (typeof data.detail === 'string') return data.detail;
if (typeof data.message === 'string') return data.message;
}
if (typeof e.message === 'string') return e.message;
}
return 'Batch fill failed.';
}
/**
* Expand role→entityIds[] into one assignment list per output variant.
@@ -383,54 +405,75 @@ export function useBatchFormFillFlow(
labelsBySourceFile[pf.fileId] = lbl;
}
// Page text per file persisted by the analyser — used to attach
// nearby_page_text to each field so the filler can disambiguate generic
// labels (e.g. "Date" near a "Contract effective date:" line).
const pageTextsByFile = analysis.pageTextsByFile ?? {};
try {
// Each variant is its own /fill-batch call: knowledge differs per assignment combo,
// and the engine treats one knowledge dict as global per request.
const filledOutputs: Array<{
plan: VariantPlan;
response: FormFillBatchResponse;
}> = [];
for (let i = 0; i < plans.length; i++) {
const plan = plans[i];
const cleanedLabels = labelsBySourceFile[plan.sourceFile.fileId] ?? {};
// Use a stable per-variant fileId so we can correlate results back.
const variantFileId = `${plan.sourceFile.fileId}__${i}`;
const batchRequest = {
files: [
{
fileId: variantFileId,
formFields: plan.fields.map((f) => ({
name: clamp(f.name, MAX_NAME) ?? '',
label: clamp(cleanedLabels[f.name] ?? f.label, MAX_LABEL),
type: f.type,
value: clamp(f.value, MAX_VALUE),
options: f.options,
displayOptions: f.displayOptions,
required: f.required,
readOnly: f.readOnly,
multiSelect: f.multiSelect,
multiline: f.multiline,
tooltip: clamp(f.tooltip, MAX_TOOLTIP),
})) as unknown as FormField[],
roleLabel: clamp(plan.assignments[0]?.roleLabel ?? 'Primary', MAX_LABEL) ?? 'Primary',
},
],
knowledge: resolveDynamicValues(plan.mergeResult.knowledge),
};
const response = await fillFormsBatch(batchRequest);
filledOutputs.push({ plan, response });
}
// Each variant is its own /fill-batch call: knowledge differs per assignment
// combo, and the engine treats one knowledge dict as global per request.
// Variants are independent, so fan them out with bounded concurrency.
const filledOutputs = await runWithConcurrency(
plans,
MAX_CONCURRENT_FILLER_CALLS,
async (plan, i): Promise<{ plan: VariantPlan; response: FormFillBatchResponse }> => {
const cleanedLabels = labelsBySourceFile[plan.sourceFile.fileId] ?? {};
const pageTexts = pageTextsByFile[plan.sourceFile.fileId] ?? {};
// Stable per-variant fileId so we can correlate results back.
const variantFileId = `${plan.sourceFile.fileId}__${i}`;
const batchRequest = {
files: [
{
fileId: variantFileId,
formFields: plan.fields.map((f) => {
const pageIndex = f.widgets?.[0]?.pageIndex;
const nearbyPageText =
pageIndex != null ? pageTexts[pageIndex] : undefined;
return {
name: clamp(f.name, MAX_NAME) ?? '',
label: clamp(cleanedLabels[f.name] ?? f.label, MAX_LABEL),
type: f.type,
value: clamp(f.value, MAX_VALUE),
options: f.options,
displayOptions: f.displayOptions,
required: f.required,
readOnly: f.readOnly,
multiSelect: f.multiSelect,
multiline: f.multiline,
tooltip: clamp(f.tooltip, MAX_TOOLTIP),
nearbyPageText: clamp(nearbyPageText, MAX_NEARBY_PAGE_TEXT),
};
}) as unknown as FormField[],
roleLabel:
clamp(plan.assignments[0]?.roleLabel ?? 'Primary', MAX_LABEL) ?? 'Primary',
},
],
knowledge: resolveDynamicValues(plan.mergeResult.knowledge),
};
const response = await fillFormsBatch(batchRequest);
return { plan, response };
},
);
// Build review proposals — engine has matched fields, user reviews/edits
// before we touch the actual PDFs.
// before we touch the actual PDFs. Failed variants (HTTP errors etc.)
// are skipped here and surfaced via the review-screen message so
// partial successes survive.
const proposed: ProposedVariant[] = [];
let emptyMatchVariants = 0;
let failedVariants = 0;
const emptyMatchEntityNames: string[] = [];
const failureMessages: string[] = [];
for (let i = 0; i < filledOutputs.length; i++) {
const { plan, response } = filledOutputs[i];
const settled = filledOutputs[i];
if (!settled.ok) {
failedVariants++;
failureMessages.push(describeError(settled.error));
continue;
}
const { plan, response } = settled.value;
const fileResult = response.perFile[0];
const filled = fileResult?.filledFields ?? [];
const cleanedLabels = labelsBySourceFile[plan.sourceFile.fileId] ?? {};
@@ -474,33 +517,56 @@ export function useBatchFormFillFlow(
});
}
// Hard-fail when nothing matched anywhere — same diagnostic as before.
if (proposed.length === 0 && emptyMatchVariants > 0) {
const uniqueEntities = Array.from(new Set(emptyMatchEntityNames));
const summary =
uniqueEntities.length === 1
? `entity "${uniqueEntities[0]}"`
: `entities ${uniqueEntities.map((n) => `"${n}"`).join(', ')}`;
const allEmpty = plans.every(
(p) =>
Object.keys(p.mergeResult.knowledge).filter((k) => !k.startsWith('_')).length === 0,
);
const reason = allEmpty
? `the assigned ${summary} ${uniqueEntities.length === 1 ? 'has' : 'have'} no fields. ` +
`Open Manage → Entities and add fields like "first_name", "address_line_1", etc.`
: `the AI couldn't match any of your entity fields to the form. ` +
`Check the field names in ${summary} — try common keys like first_name, last_name, ` +
`address_line_1, email, phone.`;
setState((s) => ({
...s,
phase: 'error',
error: `Generated 0 filled files — ${reason}`,
plannedVariantCount: plans.length,
}));
return;
// Hard-fail only when EVERY variant either errored or matched zero
// fields — otherwise partial successes proceed to review with a
// message describing what was lost.
if (proposed.length === 0) {
if (failedVariants > 0) {
// Only failures, no empty-matches: surface the underlying error.
const uniq = Array.from(new Set(failureMessages));
setState((s) => ({
...s,
phase: 'error',
error:
`All ${failedVariants} variant${failedVariants === 1 ? '' : 's'} failed: ` +
uniq.slice(0, 2).join('; ') +
(uniq.length > 2 ? ` (+${uniq.length - 2} more)` : ''),
plannedVariantCount: plans.length,
}));
return;
}
if (emptyMatchVariants > 0) {
const uniqueEntities = Array.from(new Set(emptyMatchEntityNames));
const summary =
uniqueEntities.length === 1
? `entity "${uniqueEntities[0]}"`
: `entities ${uniqueEntities.map((n) => `"${n}"`).join(', ')}`;
const allEmpty = plans.every(
(p) =>
Object.keys(p.mergeResult.knowledge).filter((k) => !k.startsWith('_')).length === 0,
);
const reason = allEmpty
? `the assigned ${summary} ${uniqueEntities.length === 1 ? 'has' : 'have'} no fields. ` +
`Open Manage → Entities and add fields like "first_name", "address_line_1", etc.`
: `the AI couldn't match any of your entity fields to the form. ` +
`Check the field names in ${summary} — try common keys like first_name, last_name, ` +
`address_line_1, email, phone.`;
setState((s) => ({
...s,
phase: 'error',
error: `Generated 0 filled files — ${reason}`,
plannedVariantCount: plans.length,
}));
return;
}
}
const reviewMessage = [
failedVariants > 0
? `${failedVariants} variant${failedVariants === 1 ? '' : 's'} failed and ${
failedVariants === 1 ? 'was' : 'were'
} skipped`
: null,
emptyMatchVariants > 0
? `${emptyMatchVariants} variant${emptyMatchVariants === 1 ? '' : 's'} skipped (no field matches)`
: null,
@@ -7,6 +7,7 @@ import { useState, useCallback } from 'react';
import { fetchFormFieldsWithCoordinates } from '@app/tools/formFill/formApi';
import { extractPageTexts } from './pdfTextExtraction';
import { analyseMultipleForms } from './aiFormFillApi';
import { describeError as describeErrorShared } from './errorUtils';
import type {
AnalysedFileResult,
CrossFileRole,
@@ -16,12 +17,19 @@ import type {
import type { FormField as FullFormField } from '@app/tools/formFill/types';
import type { KnowledgeStore } from './useKnowledgeStore';
const describeError = (err: unknown) => describeErrorShared(err, 'Form analysis failed.');
export type AnalysisPhase = 'idle' | 'fetching_fields' | 'analysing' | 'done' | 'error';
export interface FormAnalysisState {
phase: AnalysisPhase;
/** Raw fields fetched from Java backend, keyed by fileId */
fieldsByFile: Record<string, FullFormField[]>;
/**
* Page text per file, keyed by fileId then by pageIndex. Persisted so the
* fill phase can attach nearby_page_text to fields without re-running pdf.js.
*/
pageTextsByFile: Record<string, Record<number, string>>;
/** AI analysis result */
analysis: FormAnalysisResponse | null;
/**
@@ -38,6 +46,7 @@ export interface FormAnalysisState {
const INITIAL_STATE: FormAnalysisState = {
phase: 'idle',
fieldsByFile: {},
pageTextsByFile: {},
analysis: null,
roleProfileMap: {},
fileRoleOverrides: {},
@@ -59,44 +68,26 @@ const MAX_TYPE = 50;
const MAX_FILE_ID = 200;
const MAX_FILE_NAME = 500;
/**
* Beyond ~200 fields per analyse call, the smart model gets sluggish and risks
* truncating its structured output. We chunk above this threshold and merge
* roles back together client-side. Each chunk stays well under the engine's
* MAX_FIELDS_PER_FILE (engine/src/stirling/contracts/form_fill.py).
* Threshold for splitting one logical file into multiple synthetic chunks
* before the analyse call. Sonnet 4.6 has a 200K context and the new phase-1
* prompt no longer enumerates every field, so chunking is mostly a safety net
* for genuinely huge forms. 1500 stays under the engine's MAX_FIELDS_PER_FILE
* (2000) with headroom for prompt overhead. See AI_LAYERS_REVIEW.md.
*/
const FIELDS_PER_CHUNK = 200;
const FIELDS_PER_CHUNK = 1500;
/**
* Frontend-only upper bound on a single PDF — beyond this we refuse rather than
* fan the analyse step into 20+ LLM calls. Independent of the server-side
* fan analyse into many LLM calls. Independent of the server-side
* MAX_FIELDS_PER_FILE cap (which only sees one chunk at a time).
*/
const MAX_FIELDS_FOR_CHUNKED_ANALYSIS = 4000;
const MAX_FIELDS_FOR_CHUNKED_ANALYSIS = 6000;
function clamp(s: string | undefined | null, max: number): string | undefined {
if (s == null) return undefined;
return s.length > max ? s.slice(0, max) : s;
}
/** Pull the most useful detail out of an axios/fetch error. */
function describeError(err: unknown): string {
if (err && typeof err === 'object') {
const e = err as any;
const data = e.response?.data;
if (data) {
// FastAPI 422 returns {detail: [{loc, msg, type}, ...]}
if (Array.isArray(data.detail)) {
return data.detail
.map((d: any) => `${(d.loc ?? []).join('.')}: ${d.msg ?? d.type ?? ''}`)
.join('; ');
}
if (typeof data.detail === 'string') return data.detail;
if (typeof data.message === 'string') return data.message;
}
if (typeof e.message === 'string') return e.message;
}
return 'Form analysis failed.';
}
/**
* Best-effort fill-in of analysis data the LLM occasionally leaves empty.
@@ -269,6 +260,7 @@ function mergeChunkedAnalysis(
});
return {
outcome: 'form_analysis',
perFile: Object.values(perFileByLogical),
crossFileRoles,
message: analysis.message,
@@ -331,8 +323,8 @@ export function useFormAnalysis(knowledge: KnowledgeStore) {
setState((s) => ({ ...s, phase: 'analysing', fieldsByFile }));
// Step 3: Refuse hopelessly oversized files upfront — analysing 4000+ fields in
// ~200-field chunks means 20+ LLM calls; users should split the document first.
// Step 3: Refuse hopelessly oversized files upfront — beyond the cap we fan
// into multiple chunked LLM calls; users should split the document first.
const oversized = files.find(
(sf) => (fieldsByFile[sf.fileId] ?? []).length > MAX_FIELDS_FOR_CHUNKED_ANALYSIS,
);
@@ -384,7 +376,20 @@ export function useFormAnalysis(knowledge: KnowledgeStore) {
}
// Step 5: Call analyser (one request, possibly with chunked synthetic files)
const rawAnalysis = await analyseMultipleForms({ files: requestFiles });
const rawWorkflowResponse = await analyseMultipleForms({ files: requestFiles });
// The analyser can refuse via FormAnalysisAmbiguousResponse — same pattern
// as PdfEditAgent's EditClarificationRequest. Surface the LLM's reason
// instead of silently fabricating roles.
if (rawWorkflowResponse.outcome === 'form_analysis_ambiguous') {
const reason = rawWorkflowResponse.reason;
const suggestion = rawWorkflowResponse.suggestion
? ` ${rawWorkflowResponse.suggestion}`
: '';
throw new Error(`AI couldn't analyse this form: ${reason}${suggestion}`);
}
const rawAnalysis = rawWorkflowResponse;
const merged =
totalChunks > 0 ? mergeChunkedAnalysis(rawAnalysis, realFileNames) : rawAnalysis;
// Step 5b: Patch up empty / missing field assignments before anything
@@ -410,6 +415,7 @@ export function useFormAnalysis(knowledge: KnowledgeStore) {
...s,
phase: 'done',
analysis,
pageTextsByFile,
roleProfileMap: initialRoleMap,
fileRoleOverrides: {},
}));