Files
AIRegulation-DocAnalysis/backend/app/application/compliance/pipeline.py
T
2026-07-02 22:03:39 +08:00

571 lines
22 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""Compliance analysis pipeline helpers.
All functions are synchronous — call them via asyncio.to_thread() in async SSE generators.
"""
from __future__ import annotations
import asyncio
import json
import os
import re
import tempfile
from typing import TYPE_CHECKING
from loguru import logger
from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential
# Shared retry policy for LLM calls: 3 attempts, exponential back-off 14 s.
_llm_retry = retry(
stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=1, min=1, max=4),
retry=retry_if_exception_type((ValueError, TimeoutError, ConnectionError)),
reraise=True,
)
if TYPE_CHECKING:
from app.application.knowledge import KnowledgeRetrievalService
from app.domain.retrieval import RetrievedChunk
from app.domain.compliance.ports import AnalysisRecord, FindingRecord
from app.services.llm.base_client import BaseLLMClient
def _extract_json(text: str):
"""Extract JSON from LLM response, tolerating markdown wrappers."""
stripped = text.strip()
match = re.search(r"```(?:json)?\s*([\s\S]*?)```", stripped)
if match:
stripped = match.group(1).strip()
try:
return json.loads(stripped)
except json.JSONDecodeError:
pass
for pattern in (r"(\[[\s\S]*\])", r"(\{[\s\S]*\})"):
m = re.search(pattern, stripped)
if m:
try:
return json.loads(m.group(1))
except json.JSONDecodeError:
continue
raise ValueError(f"No valid JSON found in LLM response: {text[:300]}")
def extract_text_from_doc_id(doc_id: str) -> str:
"""Fetch the full text of a document by retrieving its chunks filtered by doc_id.
Uses a high top_k and doc_id filter to reconstruct the document in chunk order,
avoiding the previous approach of semantic search by doc_name which could return
chunks from unrelated documents.
"""
from app.shared.bootstrap import get_document_query_service, get_retrieval_service
doc = get_document_query_service().get(doc_id)
if not doc:
raise ValueError(f"Document '{doc_id}' not found")
service = get_retrieval_service()
# Use doc_name as a broad query, filter strictly by doc_id so we only get
# this document's chunks; top_k=100 covers most real-world documents.
chunks = service.retrieve(query=doc.doc_name, top_k=100, filters=doc_id)
doc_chunks = [c for c in chunks if getattr(c, "doc_id", None) == doc_id]
if not doc_chunks:
# Fallback: use top results even without doc_id match (e.g., legacy store)
doc_chunks = chunks[:30]
# Sort by chunk_index to preserve document reading order
doc_chunks.sort(key=lambda c: getattr(c, "chunk_index", 0))
return "\n\n".join(c.text for c in doc_chunks[:40])
def extract_text_from_file(content: bytes, filename: str) -> str:
"""Parse an uploaded file and return its full text content.
Removed previous 4000-char cap so large specifications and standards are
fully analysed. The caller is responsible for splitting the text into
clause-sized chunks before passing to the LLM.
"""
from app.shared.bootstrap import get_document_command_service
suffix = os.path.splitext(filename or "doc.pdf")[1] or ".pdf"
tmp_path = ""
try:
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
tmp.write(content)
tmp_path = tmp.name
service = get_document_command_service()
parsed = service.parser.parse(file_path=tmp_path, doc_id="tmp_analysis", doc_name=filename)
if parsed.raw_text:
# Return full text — truncation happens in split_into_clauses()
return parsed.raw_text
return "\n".join(
b.get("text", "") for b in parsed.semantic_blocks if b.get("text")
)
except Exception as exc:
logger.warning("File text extraction failed: {}", exc)
return ""
finally:
if tmp_path:
try: os.unlink(tmp_path)
except OSError: pass
def split_into_clauses(text: str, client: "BaseLLMClient") -> list[str]:
"""Split a compliance document into semantically independent clauses.
For long texts (> 2 000 chars) the document is processed in overlapping
2 000-char windows so no content is missed. Each window produces up to 4
clauses; results are deduplicated and capped at 12 total to keep analysis
latency reasonable.
"""
# Window size and step for sliding-window clause extraction
_WINDOW = 2000
_STEP = 1800 # 200-char overlap to avoid cutting clauses at boundaries
_MAX_CLAUSES = 12
windows = []
if len(text) <= _WINDOW:
windows = [text]
else:
pos = 0
while pos < len(text):
windows.append(text[pos: pos + _WINDOW])
pos += _STEP
all_clauses: list[str] = []
for window in windows:
prompt = (
"You are a compliance analysis expert. Split the following text into "
"3-4 semantically complete compliance clauses. Each clause must be an "
"independent requirement or technical statement. Omit section headings, "
"definitions, and non-normative text.\n"
"Return as JSON array of strings, e.g.:\n"
'["Clause one...", "Clause two..."]\n'
"Return ONLY the JSON array.\n\n"
f"Text:\n{window}"
)
response = client.chat([{"role": "user", "content": prompt}], max_tokens=800)
if response.is_success:
try:
result = _extract_json(response.content)
if isinstance(result, list):
clauses = [str(c).strip() for c in result if str(c).strip()]
all_clauses.extend(clauses[:4])
except (ValueError, TypeError):
logger.warning("Clause split JSON parse failed for window, using sentence fallback")
sentences = re.split(r"[.?!;\n]+", window)
all_clauses.extend(s.strip() for s in sentences if len(s.strip()) > 20)
else:
# LLM unavailable — fall back to sentence splitting for this window
sentences = re.split(r"[.?!;\n]+", window)
all_clauses.extend(s.strip() for s in sentences if len(s.strip()) > 20)
if len(all_clauses) >= _MAX_CLAUSES:
break
# Deduplicate near-duplicates (same first 80 chars) that span window boundaries
seen: set[str] = set()
deduped: list[str] = []
for c in all_clauses:
key = c[:80].lower()
if key not in seen:
seen.add(key)
deduped.append(c)
return deduped[:_MAX_CLAUSES]
def retrieve_for_clause(
clause: str,
retrieval_service: "KnowledgeRetrievalService",
top_k: int = 5,
domains: str | None = None,
) -> list["RetrievedChunk"]:
"""Retrieve regulation chunks relevant to a clause.
If the best retrieval score is below 0.55, rewrite the clause into a more
technical query and retry once to improve coverage.
"""
chunks = retrieval_service.retrieve(query=clause, top_k=top_k, filters=domains)
if not chunks:
return chunks
best_score = max((getattr(c, "score", 0) for c in chunks), default=0)
if best_score < 0.55:
# Rewrite clause as technical keyword query and retry
keywords = " ".join(
w for w in re.split(r"\W+", clause) if len(w) > 3
)[:200]
retry_chunks = retrieval_service.retrieve(query=keywords, top_k=top_k, filters=domains)
if retry_chunks:
# Merge: keep unique chunks, prefer higher-score version
seen_ids: set[str] = {getattr(c, "chunk_id", str(i)) for i, c in enumerate(chunks)}
for rc in retry_chunks:
rid = getattr(rc, "chunk_id", "")
if rid not in seen_ids:
chunks.append(rc)
seen_ids.add(rid)
chunks.sort(key=lambda c: getattr(c, "score", 0), reverse=True)
chunks = chunks[:top_k]
return chunks
def process_single_clause(
clause: str,
index: int,
retrieval_service: "KnowledgeRetrievalService",
client: "BaseLLMClient",
top_k: int = 5,
domains: str | None = None,
) -> dict:
"""Process one clause: retrieve relevant regulations then check compliance.
Returns a dict with keys:
- index: clause position (for ordering)
- chunks: list of RetrievedChunk (for source events)
- finding: dict with title/desc/status/clause_ref/confidence (may be None on LLM failure)
Designed to run inside asyncio.to_thread() for parallel execution.
The finding now includes a 'source_refs' list linking back to the chunks
that informed the verdict, enabling the frontend to correlate sources with findings.
"""
chunks = retrieve_for_clause(clause, retrieval_service, top_k, domains)
finding = check_clause_compliance(clause, chunks, client)
if finding is not None:
# Attach source references so the frontend can link finding ↔ sources
finding["source_refs"] = [
{
"standard": getattr(c, "doc_title", "") or getattr(c, "doc_name", ""),
"clause": getattr(c, "section_title", "") or "",
"score": round(float(getattr(c, "score", 0)), 3),
}
for c in chunks[:3]
]
return {"index": index, "chunks": chunks, "finding": finding}
async def run_clauses_streaming(
clauses: list[str],
retrieval_service: "KnowledgeRetrievalService",
client: "BaseLLMClient",
top_k: int = 5,
domains: str | None = None,
):
"""Process all clauses concurrently and yield each result as it completes.
Unlike the old gather()-based approach, this uses asyncio.Queue so that
findings are emitted to the SSE stream immediately when each clause
finishes — the user sees results progressively rather than waiting for
the slowest clause before seeing any output.
Yields dicts with keys: index, chunks, finding (same schema as
process_single_clause, plus a sentinel {"_done": True} at the end).
"""
queue: asyncio.Queue[dict] = asyncio.Queue()
total = len(clauses)
async def _worker(clause: str, i: int) -> None:
"""Run one clause in a thread and push the result into the queue."""
try:
result = await asyncio.to_thread(
process_single_clause,
clause, i, retrieval_service, client, top_k, domains,
)
except Exception as exc:
logger.warning("Clause {} processing failed: {}", i, exc)
result = {"index": i, "chunks": [], "finding": None}
await queue.put(result)
# Launch all workers concurrently
tasks = [asyncio.create_task(_worker(clause, i)) for i, clause in enumerate(clauses)]
received = 0
while received < total:
result = await queue.get()
yield result
received += 1
# Wait for all tasks to complete (they should already be done by now)
await asyncio.gather(*tasks, return_exceptions=True)
async def run_clauses_parallel(
clauses: list[str],
retrieval_service: "KnowledgeRetrievalService",
client: "BaseLLMClient",
top_k: int = 5,
domains: str | None = None,
) -> list[dict]:
"""Legacy batch API kept for backward compatibility.
Collects all streaming results and returns them sorted by clause index.
New code should use run_clauses_streaming() directly.
"""
results: list[dict] = []
async for result in run_clauses_streaming(clauses, retrieval_service, client, top_k, domains):
results.append(result)
return sorted(results, key=lambda r: r["index"])
def check_clause_compliance(
clause: str,
chunks: list["RetrievedChunk"],
client: "BaseLLMClient",
) -> dict | None:
"""Check whether a business clause complies with the retrieved regulations.
The prompt explicitly instructs the LLM to:
- extract clause_ref from the retrieved text (not invent it)
- include a confidence score (0-1) reflecting how well the retrieved
chunks cover the clause topic
Returns None only when the LLM call fails after all retries.
"""
reg_context = "\n".join(
f"[{i+1}] {c.doc_title} {c.section_title or ''}: {c.text[:300]}"
for i, c in enumerate(chunks[:5])
) if chunks else "(no regulatory context retrieved)"
prompt = (
"You are a compliance expert. Judge whether the following business clause "
"complies with the retrieved regulations.\n\n"
f"Business clause:\n{clause}\n\n"
f"Retrieved regulations:\n{reg_context}\n\n"
"Return JSON with these exact fields:\n"
"{\n"
' "status": "ok" | "warn" | "risk",\n'
' "title": "Short finding title (max 30 chars)",\n'
' "desc": "Description (50-120 chars)",\n'
' "clause_ref": "Exact clause/article reference copied from the retrieved text above, '
'e.g. Art.9.1 or Sec.3.1. Use null if no specific clause number appears in the retrieved text.",\n'
' "confidence": 0.0-1.0 // how well the retrieved context covers this clause topic\n'
"}\n"
"status: ok=compliant, warn=gap exists, risk=critical/missing\n"
"IMPORTANT: copy clause_ref verbatim from the retrieved text; do NOT invent references.\n"
"Return ONLY the JSON object."
)
def _do_check():
resp = client.chat([{"role": "user", "content": prompt}], max_tokens=500)
if not resp.is_success:
raise ValueError("LLM returned non-success for gap check")
return resp
try:
response = _llm_retry(_do_check)()
except Exception as exc:
logger.warning("check_clause_compliance LLM call failed after retries: {}", exc)
return None
try:
result = _extract_json(response.content)
if isinstance(result, dict) and "status" in result:
return {
"title": str(result.get("title", "Compliance finding")),
"desc": str(result.get("desc", "")),
"status": result.get("status", "info"),
# None if LLM correctly found no clause number in retrieved text
"clause_ref": result.get("clause_ref") or None,
# Confidence score helps frontend show retrieval quality indicator
"confidence": float(result.get("confidence", 0.5)),
}
except (ValueError, TypeError) as exc:
logger.warning("Gap check JSON parse failed: {}", exc)
return None
def synthesize_conclusion(
para_text: str,
findings: list[dict],
client: "BaseLLMClient",
) -> dict:
if not findings:
return {
"conclusion": "No significant compliance gaps found. Continue monitoring regulation updates.",
"actions": [{"label": "Next action", "value": "Monitor regulation updates"}],
"risk_score": 10,
"highlight_terms": [],
"para_text": para_text[:800],
}
findings_text = "\n".join(
f"- [{f['status'].upper()}] {f['title']}: {f['desc']}"
for f in findings
)
prompt = (
"You are a compliance analysis expert. Generate a summary report "
"based on the following compliance findings.\n\n"
f"Original text (first 600 chars):\n{para_text[:600]}\n\n"
f"Findings:\n{findings_text}\n\n"
"Return JSON:\n"
"{\n"
' "conclusion": "Overall compliance conclusion (100-200 chars)",\n'
' "actions": [\n'
' {"label": "Action label", "value": "Description"},\n'
' {"label": "Priority", "value": "High/Medium/Low", "risk": true}\n'
' ],\n'
' "risk_score": 0-100 (integer, higher=riskier),\n'
' "highlight_terms": ["term1", "term2"], // up to 10 key technical/legal terms actually present in the text\n'
' "para_text": "Original text or summary (max 600 chars)"\n'
"}\n"
"Return ONLY the JSON object."
)
fallback = {
"conclusion": "Compliance analysis complete. Review findings and create remediation plan.",
"actions": [
{"label": "Next action", "value": "Review critical findings"},
{"label": "Escalation", "value": "Legal review required", "risk": True},
],
"risk_score": 60,
"highlight_terms": [],
"para_text": para_text[:800],
}
def _do_synthesize():
resp = client.chat([{"role": "user", "content": prompt}], max_tokens=1200)
if not resp.is_success:
raise ValueError("LLM returned non-success for synthesis")
return resp
try:
response = _llm_retry(_do_synthesize)()
except Exception as exc:
logger.warning("synthesize_conclusion LLM call failed after retries: {}", exc)
return fallback
try:
result = _extract_json(response.content)
if isinstance(result, dict):
return {
"conclusion": str(result.get("conclusion", fallback["conclusion"])),
"actions": result.get("actions", fallback["actions"]),
"risk_score": int(result.get("risk_score", 60)),
"highlight_terms": result.get("highlight_terms", []),
"para_text": str(result.get("para_text", para_text[:800])),
}
except (ValueError, TypeError) as exc:
logger.warning("Conclusion synthesis JSON parse failed: {}", exc)
return fallback
_SUGGESTION_FOCUS = {
"risk": "Focus on remediation steps, required certifications, and timeline to resolve.",
"warn": "Focus on identifying the specific compliance gap and how to close it.",
"ok": "Focus on maintaining compliance evidence and monitoring future changes.",
}
_SUGGESTION_FALLBACK = {
"risk": [
"What specific certifications or documents are required to remediate this finding?",
"What is the typical remediation timeline for this type of non-compliance?",
"Which regulation clause defines the exact requirement?",
],
"warn": [
"What is the exact gap between the current state and the requirement?",
"What evidence would demonstrate partial compliance?",
"Which regulation clause applies to this warning?",
],
"ok": [
"What documentation should be maintained to evidence this compliance?",
"How should this area be monitored as regulations evolve?",
"Are there related clauses that may affect this compliant area?",
],
}
def build_finding_context(finding: "FindingRecord", analysis: "AnalysisRecord") -> str:
"""Build a grounded system context string for a finding chat thread.
Combines finding details with analysis metadata so the LLM has full
context without relying on the frontend to pass segment_context.
"""
return (
f"Document: {analysis.doc_name}\n"
f"Standard: {analysis.standard_name}\n"
f"Finding [{finding.seq + 1}]: {finding.title}\n"
f"Status: {finding.status}\n"
f"Clause reference: {finding.clause_ref or 'N/A'}\n"
f"Description: {finding.description}\n"
f"Overall conclusion: {analysis.conclusion}\n"
)
def generate_suggestions(
finding: "FindingRecord",
analysis: "AnalysisRecord",
client: "BaseLLMClient",
) -> list[str]:
"""Generate 3 context-aware follow-up questions for a finding chat thread.
Returns exactly 3 question strings. Falls back to static templates on error.
"""
fallback = _SUGGESTION_FALLBACK.get(finding.status, _SUGGESTION_FALLBACK["warn"])
context = build_finding_context(finding, analysis)
focus = _SUGGESTION_FOCUS.get(finding.status, _SUGGESTION_FOCUS["warn"])
prompt = (
f"{context}\n\n"
f"Task: {focus}\n"
"Generate exactly 3 concise follow-up questions a compliance analyst would ask.\n"
'Return JSON: {"questions": ["question 1", "question 2", "question 3"]}\n'
"Return ONLY the JSON object."
)
response = client.chat([{"role": "user", "content": prompt}], max_tokens=300)
if not response.is_success:
return fallback
try:
result = _extract_json(response.content)
questions = result.get("questions", [])
if isinstance(questions, list) and len(questions) >= 3:
return [str(q) for q in questions[:3]]
except (ValueError, TypeError) as exc:
logger.warning("generate_suggestions JSON parse failed: {}", exc)
return fallback
def detect_cross_clause_conflicts(
findings: list[dict],
client: "BaseLLMClient",
) -> list[dict]:
"""Detect contradictions and missing cross-references across all findings.
Runs a single LLM call after all per-clause findings are collected.
Returns a list of conflict dicts: {type, finding_a, finding_b, desc}.
Returns an empty list on LLM failure so the caller can proceed without it.
"""
if len(findings) < 2:
# Need at least 2 findings to compare
return []
findings_text = "\n".join(
f"[{i+1}] [{f['status'].upper()}] {f['title']}: {f['desc']}"
+ (f" (Ref: {f['clause_ref']})" if f.get("clause_ref") else "")
for i, f in enumerate(findings)
)
prompt = (
"You are a compliance expert. Review the following compliance findings from the same document "
"and identify any cross-clause issues:\n\n"
f"Findings:\n{findings_text}\n\n"
"Return JSON array of conflicts (empty array [] if none found):\n"
"[\n"
" {\n"
' "type": "contradiction" | "missing_ref" | "cumulative_risk",\n'
' "finding_a": <1-based index>,\n'
' "finding_b": <1-based index or null>,\n'
' "desc": "Brief description of the cross-clause issue (max 100 chars)"\n'
" }\n"
"]\n"
"Return ONLY the JSON array."
)
try:
response = client.chat([{"role": "user", "content": prompt}], max_tokens=600)
if not response.is_success:
return []
result = _extract_json(response.content)
if isinstance(result, list):
return [
{
"type": str(c.get("type", "contradiction")),
"finding_a": int(c.get("finding_a", 0)),
"finding_b": c.get("finding_b"),
"desc": str(c.get("desc", "")),
}
for c in result
if isinstance(c, dict)
]
except Exception as exc:
logger.warning("detect_cross_clause_conflicts failed: {}", exc)
return []