"""LangGraph extract 그래프: extract → validate → (retry) → normalize.""" from __future__ import annotations import base64 from typing import Any from langgraph.graph import END, StateGraph from pydantic import ValidationError from app.config import settings from app.graph.prompts import ( EXTRACTION_FEWSHOT, EXTRACTION_SYSTEM_PROMPT, RETRY_SYSTEM_PROMPT, ) from app.graph.state import ExtractState from app.models.ir import COMPONENT_TYPES, Background, Layout, ReportIR from app.services import llm as llm_service def extract_node(state: ExtractState) -> ExtractState: image_bytes = base64.b64decode(state["image_b64"]) raw = llm_service.call_vision( EXTRACTION_SYSTEM_PROMPT, EXTRACTION_FEWSHOT, image_bytes, state.get("mime", "image/png") ) return {"raw_output": raw, "attempts": state.get("attempts", 0) + 1} def _validate(raw: str) -> tuple[dict[str, Any] | None, list[str]]: try: parsed = llm_service.parse_json(raw) except Exception as e: # noqa: BLE001 return None, [f"JSON 파싱 실패: {e}"] try: ir = ReportIR.model_validate(parsed) except ValidationError as e: return None, [str(err) for err in e.errors()][:10] return ir.dump(), [] def validate_node(state: ExtractState) -> ExtractState: ir, errors = _validate(state.get("raw_output", "")) if ir is not None: return {"ir": ir, "errors": []} return {"errors": errors} def retry_node(state: ExtractState) -> ExtractState: prompt = RETRY_SYSTEM_PROMPT.format( errors="\n".join(state.get("errors", [])), previous=state.get("raw_output", "")[:4000], ) image_bytes = base64.b64decode(state["image_b64"]) raw = llm_service.call_vision(prompt, "", image_bytes, state.get("mime", "image/png")) return {"raw_output": raw, "attempts": state.get("attempts", 0) + 1} _DATA_TYPES = { "line_chart", "bar_chart", "pie_chart", "metric_table", "holdings_table", "table", "bullet_list", "paragraph", } def _heading_text(comp: dict[str, Any]) -> str: if comp.get("title"): return comp["title"] for f in comp.get("dataSpec", {}).get("fields", []) or []: if f.get("label"): return f["label"] return "" def _x_overlap(a: dict[str, Any], b: dict[str, Any]) -> bool: ax, aw = a.get("x", 0), a.get("w", 12) bx, bw = b.get("x", 0), b.get("w", 12) return ax < bx + bw and bx < ax + aw def _merge_section_headings(comps: list[dict[str, Any]]) -> list[dict[str, Any]]: """제목(section_heading)을 바로 아래 인접한 데이터 컴포넌트의 title 로 흡수.""" removed: set[int] = set() for i, head in enumerate(comps): if head.get("type") != "section_heading": continue hl = head.get("layout", {}) best, best_dy = None, None for j, cand in enumerate(comps): if j == i or j in removed or cand.get("type") not in _DATA_TYPES: continue cl = cand.get("layout", {}) dy = cl.get("y", 0) - hl.get("y", 0) if dy < 0 or not _x_overlap(hl, cl): continue if best_dy is None or dy < best_dy: best, best_dy = cand, dy if best is not None and best_dy is not None and best_dy <= 3: text = _heading_text(head) if text: best["title"] = text bl = best["layout"] new_y = min(bl.get("y", 0), hl.get("y", 0)) bl["h"] = bl.get("h", 1) + (bl.get("y", 0) - new_y) bl["y"] = new_y hb, cb = head.get("bbox"), best.get("bbox") if hb and cb: top = min(hb["y"], cb["y"]) left = min(hb["x"], cb["x"]) right = max(hb["x"] + hb["w"], cb["x"] + cb["w"]) bottom = max(hb["y"] + hb["h"], cb["y"] + cb["h"]) best["bbox"] = {"x": left, "y": top, "w": right - left, "h": bottom - top} removed.add(i) return [c for k, c in enumerate(comps) if k not in removed] def _normalize_bbox(bbox: dict[str, Any] | None, layout: dict[str, Any]) -> dict[str, float]: if isinstance(bbox, dict) and all(k in bbox for k in ("x", "y", "w", "h")): x = min(max(float(bbox["x"]), 0.0), 1.0) y = min(max(float(bbox["y"]), 0.0), 1.0) w = min(max(float(bbox["w"]), 0.01), 1.0 - x) h = min(max(float(bbox["h"]), 0.01), 1.0 - y) return {"x": round(x, 4), "y": round(y, 4), "w": round(w, 4), "h": round(h, 4)} rows = 50.0 lx, ly = layout.get("x", 0), layout.get("y", 0) lw, lh = layout.get("w", 12), layout.get("h", 1) x = min(max(lx / 12.0, 0.0), 1.0) y = min(max(ly / rows, 0.0), 1.0) return {"x": round(x, 4), "y": round(y, 4), "w": round(min(lw / 12.0, 1.0 - x), 4), "h": round(min(lh / rows, 1.0 - y), 4)} def normalize_node(state: ExtractState) -> ExtractState: ir = state.get("ir") or {"schemaVersion": "1.0", "report": {}} report = ir.setdefault("report", {}) bg = report.get("background") if not bg or bg.get("type") != "background_image": report["background"] = Background().model_dump(by_alias=True) report.setdefault("grid", {"cols": 12, "rowUnit": 33, "gap": 8}) comps = [c for c in (report.get("components", []) or []) if c.get("type") != "background_image"] for comp in comps: if comp.get("type") not in COMPONENT_TYPES: comp["type"] = "paragraph" comp["confidence"] = min(comp.get("confidence", 0.3), 0.3) layout = comp.get("layout") or {} x = int(layout.get("x", 0)); w = int(layout.get("w", 12)) if x < 0: x = 0 if x + w > 12: w = max(1, 12 - x) comp["layout"] = {"x": x, "y": int(layout.get("y", 0)), "w": w, "h": int(layout.get("h", 1))} comp["bbox"] = _normalize_bbox(comp.get("bbox"), comp["layout"]) comp.setdefault("dataSpec", {"fields": []}) comp["dataSpec"].setdefault("fields", []) comps = _merge_section_headings(comps) confidences: list[float] = [] for idx, comp in enumerate(comps): comp["id"] = f"cmp_{idx:03d}" confidences.append(float(comp.get("confidence", 1.0))) report["components"] = comps report["overallConfidence"] = round(sum(confidences) / len(confidences), 3) if confidences else 0.0 ir["report"] = report return {"ir": ir, "confidence": report["overallConfidence"]} def _route_after_validate(state: ExtractState) -> str: if state.get("ir"): return "normalize" if state.get("attempts", 0) > settings.llm_max_retries: # 최종 실패 → 빈 IR로 normalize (배경만이라도 생성) return "normalize" return "retry" def build_extract_graph(): g = StateGraph(ExtractState) g.add_node("extract", extract_node) g.add_node("validate", validate_node) g.add_node("retry", retry_node) g.add_node("normalize", normalize_node) g.set_entry_point("extract") g.add_edge("extract", "validate") g.add_conditional_edges("validate", _route_after_validate, {"retry": "retry", "normalize": "normalize"}) g.add_edge("retry", "validate") g.add_edge("normalize", END) return g.compile()