Spaces:
Sleeping
Sleeping
File size: 4,922 Bytes
db4ba8d 7a16f4e db4ba8d 1cf88ff db4ba8d dd9584b db4ba8d 7a16f4e db4ba8d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | """
TradeFlow AI — LangGraph Extraction Graph (Step 2 Assembly)
PRD §10 — Full LangGraph pipeline:
preprocess → llm_extraction → [fallback if needed] → validate
→ risk_assessment → [interrupt if review needed] → DONE
The graph is compiled with a Redis checkpointer for persistence
and resumability across server restarts.
"""
from __future__ import annotations
import redis
import structlog
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import END, StateGraph
from ..config import settings
from .nodes.extract import llm_extraction_node
from .nodes.fallback_ocr import fallback_ocr_node
from .nodes.human_review import human_review_node
from .nodes.preprocess import preprocess_documents_node
from .nodes.risk import risk_assessment_node
from .nodes.validate import validation_node
from .state import ExtractionGraphState
log = structlog.get_logger()
def _needs_fallback(state: ExtractionGraphState) -> str:
"""Route to OCR ensemble fallback when quality, confidence, or data is weak."""
if settings.CLOUD_LLM_ONLY:
log.info("CLOUD_LLM_ONLY is active — bypassing heavy OCR ensemble fallback")
return "validate"
for doc in state.get("documents", []):
if doc.get("error") or not doc.get("extracted_data"):
return "fallback"
if doc.get("quality_score", 1.0) < settings.OCR_FALLBACK_TRIGGER_QUALITY:
return "fallback"
if doc.get("document_mode") == "digital_pdf_text" and doc.get("ocr_method") == "digital_text_parser":
continue
confidences = doc.get("field_confidences") or {}
if confidences and min(confidences.values()) < settings.OCR_FALLBACK_TRIGGER_CONFIDENCE:
return "fallback"
if doc.get("ocr_conflicts"):
return "fallback"
if len(doc.get("ocr_candidates") or {}) > 1 and doc.get("document_mode") != "digital_pdf_text":
return "fallback"
return "validate"
def _needs_review(state: ExtractionGraphState) -> str:
"""Conditional edge: route to human review if flagged."""
if state.get("needs_human_review", False):
return "human_review"
return END
def build_extraction_graph() -> StateGraph:
"""Build and compile the LangGraph extraction pipeline."""
workflow = StateGraph(ExtractionGraphState)
# ── Add nodes ────────────────────────────────────────────────
workflow.add_node("preprocess", preprocess_documents_node)
workflow.add_node("llm_extraction", llm_extraction_node)
workflow.add_node("fallback_ocr", fallback_ocr_node)
workflow.add_node("validate", validation_node)
workflow.add_node("risk_assessment", risk_assessment_node)
workflow.add_node("human_review", human_review_node)
# ── Entry point ───────────────────────────────────────────────
workflow.set_entry_point("preprocess")
# ── Edges ─────────────────────────────────────────────────────
workflow.add_edge("preprocess", "llm_extraction")
# After extraction: check if fallback needed
workflow.add_conditional_edges(
"llm_extraction",
_needs_fallback,
{
"fallback": "fallback_ocr",
"validate": "validate",
},
)
# Fallback always proceeds to validate
workflow.add_edge("fallback_ocr", "validate")
# After validation: compute risk
workflow.add_edge("validate", "risk_assessment")
# After risk: check if human review needed
workflow.add_conditional_edges(
"risk_assessment",
_needs_review,
{
"human_review": "human_review",
END: END,
},
)
# After human review: graph ends (operator approved)
workflow.add_edge("human_review", END)
return workflow
def get_compiled_graph():
"""
Returns the compiled graph with Redis checkpointer.
The checkpointer enables:
- State persistence across Celery task restarts
- interrupt() resumability for human review
- LangSmith tracing integration
"""
workflow = build_extraction_graph()
# Using MemorySaver to support async ainvoke correctly
checkpointer = MemorySaver()
graph = workflow.compile(checkpointer=checkpointer, interrupt_before=["human_review"])
log.info("LangGraph extraction graph compiled", nodes=list(workflow.nodes.keys()))
return graph
# ── Singleton (imported by Celery tasks) ──────────────────────────────────────
extraction_graph = get_compiled_graph()
|