Update core/extract.py

core/extract.py

@@ -1,9 +1,60 @@
+# core/extract.py
 from __future__ import annotations
-import json
-from typing import Any, Dict, List, Optional
-from .openai_client import get_client, VISION_MODEL, TEXT_MODEL
+import os, io, base64, json, shutil
+from typing import List, Tuple, Dict, Any
 
-SYSTEM_JSON = """あなたは有能な財務アナリストです。
+import pandas as pd
+import pdfplumber
+from pdf2image import convert_from_path
+from openai import OpenAI
+
+# ==== モデル指定（環境変数で変更可） ====
+OPENAI_MODEL_VISION = os.environ.get("OPENAI_VISION_MODEL", "gpt-4o-mini")
+OPENAI_MODEL_TEXT   = os.environ.get("OPENAI_TEXT_MODEL",   "gpt-4o-mini")
+
+# ==== 共通ユーティリティ ====
+def _b64(b: bytes) -> str:
+    return base64.b64encode(b).decode("utf-8")
+
+def _client() -> OpenAI:
+    """
+    OpenAI公式 SDK v1 系。httpx との互換のため requirements は httpx==0.27.* を推奨。
+    """
+    key = os.environ.get("OPENAI_API_KEY")
+    if not key:
+        raise RuntimeError("OPENAI_API_KEY が未設定です（Spaces → Settings → Variables and secrets）。")
+    # proxies は渡さない（互換性エラーを避ける）
+    return OpenAI(api_key=key, timeout=30)
+
+# ==== PDF 読み込み ====
+def _pdf_to_images(path: str, dpi: int = 220, max_pages: int = 6) -> List[bytes]:
+    """
+    Poppler系バイナリ（pdftoppm/pdftocairo）が必要です。Spaces なら packages.txt に
+    `poppler-utils` を入れておくと安定します。
+    """
+    imgs: List[bytes] = []
+    pages = convert_from_path(path, dpi=dpi, fmt="png")
+    for i, p in enumerate(pages):
+        if i >= max_pages:
+            break
+        buf = io.BytesIO()
+        p.save(buf, format="PNG")
+        imgs.append(buf.getvalue())
+    return imgs
+
+def _pdf_to_text(path: str, max_chars: int = 15000) -> str:
+    out: List[str] = []
+    with pdfplumber.open(path) as pdf:
+        for i, page in enumerate(pdf.pages):
+            t = (page.extract_text() or "").strip()
+            if t:
+                out.append(f"[page {i+1}]\n{t}")
+            if sum(len(x) for x in out) > max_chars:
+                break
+    return "\n\n".join(out)[:max_chars]
+
+# ==== LLM へ渡す JSON 指示 ====
+_SYSTEM_JSON = """あなたは有能な財務アナリストです。
 与えられた決算書（画像またはテキスト）から、次の厳密な JSON 構造のみを日本語の単位なし・半角数値で返してください。分からない項目は null。
 {
   "company": {"name": null},
@@ -24,42 +75,104 @@ SYSTEM_JSON = """あなたは有能な財務アナリストです。
 }
 """
 
-def extract_financials(images: Optional[List[bytes]], text_blob: Optional[str], company_hint: str="") -> Dict[str, Any]:
-    client = get_client()
+def _extract_with_llm(images: List[bytes] | None, text_blob: str | None, company_hint: str) -> Dict[str, Any]:
+    client = _client()
     if images:
-        content = [{"type": "text", "text": SYSTEM_JSON}]
+        content = [{"type": "text", "text": _SYSTEM_JSON}]
         if company_hint:
             content.append({"type": "text", "text": f"会社名の候補: {company_hint}"})
         for im in images:
-            content.append({"type": "input_image", "image_url": f"data:image/png;base64,{im.decode('latin1') if isinstance(im, str) else 'data'}"})
-        # 上のデータ URI 生成は UI 側で行うためここでは未使用
-        # （UIでdata:image/png;base64,xxxを組む実装に合わせる場合は差し替え）
-        pass
-    # 実運用では UI 側で Vision を呼ぶ形にせず、ここで共通化
-    if images:
-        content = [{"type":"text","text":SYSTEM_JSON}]
-        for im in images:
-            import base64
-            content.append({"type":"input_image","image_url":f"data:image/png;base64,{base64.b64encode(im).decode('utf-8')}"})
+            content.append({"type": "input_image", "image_url": f"data:image/png;base64,{_b64(im)}"})
         resp = client.chat.completions.create(
-            model=VISION_MODEL,
+            model=OPENAI_MODEL_VISION,
             messages=[
-                {"role":"system","content":"返答は必ず有効な JSON オブジェクトのみ。説明は不要。"},
-                {"role":"user","content":content},
+                {"role": "system", "content": "返答は必ず有効な JSON オブジェクトのみ。説明を含めない。"},
+                {"role": "user", "content": content},
             ],
-            response_format={"type":"json_object"},
+            response_format={"type": "json_object"},
             temperature=0.1,
         )
         return json.loads(resp.choices[0].message.content)
     else:
-        prompt = f"{SYSTEM_JSON}\n\n以下は決算書のテキストです。上記の JSON だけを返してください。\n\n{text_blob or ''}"
+        prompt = f"{_SYSTEM_JSON}\n\n以下は決算書のテキストです。上記の JSON だけを返してください。\n\n{text_blob or ''}"
         resp = client.chat.completions.create(
-            model=TEXT_MODEL,
+            model=OPENAI_MODEL_TEXT,
             messages=[
-                {"role":"system","content":"返答は必ず有効な JSON オブジェクトのみ。説明は不要。"},
-                {"role":"user","content":prompt},
+                {"role": "system", "content": "返答は必ず有効な JSON オブジェクトのみ。"},
+                {"role": "user", "content": prompt},
             ],
-            response_format={"type":"json_object"},
+            response_format={"type": "json_object"},
             temperature=0.1,
         )
         return json.loads(resp.choices[0].message.content)
+
+# ==== JSON <-> DataFrame ====
+def _fin_to_df(fin: Dict[str, Any]) -> pd.DataFrame:
+    rows: List[Dict[str, Any]] = []
+    def add(cat: str, d: Dict[str, Any] | None):
+        for k, v in (d or {}).items():
+            rows.append({"category": cat, "item": k, "value": v})
+    add("balance_sheet", fin.get("balance_sheet"))
+    add("income_statement", fin.get("income_statement"))
+    add("cash_flows", fin.get("cash_flows"))
+    return pd.DataFrame(rows, columns=["category", "item", "value"])
+
+# ==== 公開 API：parse_pdf ====
+def parse_pdf(files: List[str], company: str = "", force_ocr: bool = False) -> Tuple[Dict[str, Any], pd.DataFrame]:
+    """
+    入力: PDFファイルパスの配列
+    出力: (抽出JSON辞書, 表編集用DataFrame)
+    方針:
+      - まず PDF→画像化して Vision で抽出（poppler が無い/失敗なら例外）
+      - 画像抽出が失敗したらテキスト抽出→Textモデルで抽出
+      - `force_ocr=True` の場合は常に画像→Vision を試みる
+    """
+    if not files:
+        raise ValueError("PDF が指定されていません。")
+
+    # 1) 画像化（複数PDFを順に）
+    images: List[bytes] = []
+    if force_ocr:
+        for p in files:
+            images += _pdf_to_images(p, dpi=220, max_pages=6)
+    else:
+        # 画像化を試して、ダメならテキストにフォールバック
+        try:
+            for p in files:
+                images += _pdf_to_images(p, dpi=220, max_pages=6)
+        except Exception:
+            images = []
+
+    # 2) Vision / Text のいずれかで抽出
+    try:
+        if images:
+            fin = _extract_with_llm(images, None, company or "")
+        else:
+            # テキスト抽出
+            text_blob = ""
+            for p in files:
+                text_blob += _pdf_to_text(p) + "\n\n"
+            fin = _extract_with_llm(None, text_blob, company or "")
+    except Exception as e:
+        # LLM失敗時も最後にテキスト抽出で最低限の骨格を返す
+        text_blob = ""
+        for p in files:
+            try:
+                text_blob += _pdf_to_text(p) + "\n\n"
+            except Exception:
+                pass
+        fin = {
+            "company": {"name": company or None},
+            "period": {"start_date": None, "end_date": None},
+            "balance_sheet": {"total_assets": None, "total_liabilities": None, "total_equity": None,
+                              "current_assets": None, "fixed_assets": None,
+                              "current_liabilities": None, "long_term_liabilities": None},
+            "income_statement": {"sales": None, "cost_of_sales": None, "gross_profit": None,
+                                 "operating_expenses": None, "operating_income": None,
+                                 "ordinary_income": None, "net_income": None},
+            "cash_flows": {"operating_cash_flow": None, "investing_cash_flow": None, "financing_cash_flow": None},
+            "_fallback_note": f"LLM抽出に失敗したため簡易骨格のみ返却（理由: {type(e).__name__})"
+        }
+
+    df = _fin_to_df(fin)
+    return fin, df