Create extract.py

core/extract.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+import json
+from typing import Any, Dict, List, Optional
+from .openai_client import get_client, VISION_MODEL, TEXT_MODEL
+from .pdf_io import b64
+
+SYSTEM_JSON = """あなたは有能な財務アナリストです。
+与えられた決算書（画像またはテキスト）から、次の厳密な JSON 構造のみを日本語の単位なし・半角数値で返してください。分からない項目は null。
+{
+  "company": {"name": null},
+  "period": {"start_date": null, "end_date": null},
+  "balance_sheet": {
+    "total_assets": null, "total_liabilities": null, "total_equity": null,
+    "current_assets": null, "fixed_assets": null,
+    "current_liabilities": null, "long_term_liabilities": null
+  },
+  "income_statement": {
+    "sales": null, "cost_of_sales": null, "gross_profit": null,
+    "operating_expenses": null, "operating_income": null,
+    "ordinary_income": null, "net_income": null
+  },
+  "cash_flows": {
+    "operating_cash_flow": null, "investing_cash_flow": null, "financing_cash_flow": null
+  }
+}
+"""
+
+def extract_financials(images: Optional[List[bytes]], text_blob: Optional[str], company_hint: str="") -> Dict[str, Any]:
+    client = get_client()
+    if images and len(images) > 0:
+        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,{b64(im)}"})
+        resp = client.chat.completions.create(
+            model=VISION_MODEL,
+            messages=[
+                {"role": "system", "content": "返答は必ず有効な JSON オブジェクトのみ。説明は不要。"},
+                {"role": "user", "content": content},
+            ],
+            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 ''}"
+        resp = client.chat.completions.create(
+            model=TEXT_MODEL,
+            messages=[
+                {"role": "system", "content": "返答は必ず有効な JSON オブジェクトのみ。説明は不要。"},
+                {"role": "user", "content": prompt},
+            ],
+            response_format={"type": "json_object"},
+            temperature=0.1,
+        )
+        return json.loads(resp.choices[0].message.content)