Update rag/ingest.py

rag/ingest.py

@@ -1,7 +1,192 @@
-# 既存の import 群に追加は不要。下記の関数をファイル末尾あたりに追記してください。
+# rag/ingest.py
+"""
+EDINET/アップロードPDFの取り込み（RAGインデックス投入）ユーティリティ
 
-def ingest_pdf_bytes(title: str, source_url: str, pdf_bytes: bytes):
-    """ローカル/アップロードPDFをインデックスへ投入"""
+- ingest_edinet_for_company(edinet_code, date)
+    指定日の EDINET 公開一覧から該当企業の docID を探し、PDFを取得してインデックスへ投入
+- ingest_pdf_bytes(title, source_url, pdf_bytes)
+    アップロード等のPDFバイト列をそのまま投入
+- download_edinet_pdf(doc_id)
+    EDINET API から docID の PDF(=type=1) をダウンロード
+
+依存:
+- irpr.deps.add_to_index  …… ベクトル化＋永続化
+- pymupdf (fitz)          …… PDF → テキスト抽出
+- requests                …… EDINET API 呼び出し
+"""
+
+from __future__ import annotations
+
+import io
+import os
+import re
+import time
+import json
+import math
+import datetime as dt
+from typing import List, Dict, Optional
+
+import requests
+import fitz  # PyMuPDF
+
+from irpr.deps import add_to_index
+
+# =============================================================================
+# EDINET API
+# =============================================================================
+
+EDINET_API_LIST = "https://disclosure.edinet-fsa.go.jp/api/v2/documents.json"
+EDINET_API_DOC = "https://disclosure.edinet-fsa.go.jp/api/v2/documents/{doc_id}"
+
+# 代表的なドキュメント種別（必要に応じて追加）
+# 120: 有価証券報告書, 130: 四半期報告書, 140: 半期報告書, 150: 臨時報告書
+# 160: 参照書類, 170: 訂正有価証券報告書, 180: 訂正四半期報告書, etc.
+PREFERRED_DOC_TYPES = {
+    "120", "130", "140", "150", "170", "180", "350", "360", "370", "380"
+}
+
+
+def _http_get(url: str, *, headers: Optional[dict] = None, params: Optional[dict] = None,
+              timeout: int = 60) -> requests.Response:
+    """GET with simple retry."""
+    last = None
+    for _ in range(3):
+        try:
+            r = requests.get(url, headers=headers, params=params, timeout=timeout)
+            if r.status_code >= 500:
+                time.sleep(1.2)
+                last = r
+                continue
+            return r
+        except Exception as e:
+            last = e
+            time.sleep(1.2)
+    if isinstance(last, requests.Response):
+        return last
+    raise RuntimeError(f"GET failed for {url}: {last!r}")
+
+
+def list_edinet_docs_for_date(date: str) -> List[dict]:
+    """
+    指定日の EDINET 公開一覧（JSON）を取得。
+    API仕様上、日付は必須・1日単位（時刻は不可）。
+    """
+    params = {"date": date, "type": 2}  # type=2 → JSON（一覧）
+    r = _http_get(EDINET_API_LIST, params=params)
+    if r.status_code != 200:
+        raise RuntimeError(f"EDINET list error: {r.status_code} {r.text[:200]}")
+    try:
+        data = r.json()
+    except json.JSONDecodeError:
+        # まれにCSVが返るケースを回避（公式はJSONだが保険）
+        raise RuntimeError("EDINET list: JSON decode error")
+    return data.get("results", []) or []
+
+
+def find_company_doc_ids(edinet_code: str, date: str, search_window_days: int = 2) -> List[dict]:
+    """
+    指定日の前後 search_window_days 日をゆるく探索し、
+    edinetCode が一致する結果を返す（スコア順：優先種別→新しい日付）。
+    """
+    base = dt.datetime.strptime(date, "%Y-%m-%d").date()
+    candidates: List[dict] = []
+
+    for off in range(-search_window_days, search_window_days + 1):
+        d = (base + dt.timedelta(days=off)).isoformat()
+        try:
+            rows = list_edinet_docs_for_date(d)
+        except Exception:
+            continue
+        for row in rows:
+            if (row.get("edinetCode") or "").upper() == edinet_code.upper():
+                # スコアリング：好ましいdocTypeを優先
+                dtc = str(row.get("docTypeCode") or "")
+                score = 100 if dtc in PREFERRED_DOC_TYPES else 0
+                # さらに当日日付に近いほど加点（0日差=+20, 1日差=+15, …）
+                score += max(0, 20 - 5 * abs(off))
+                row["_score"] = score
+                row["_date"] = d
+                candidates.append(row)
+
+    # スコア降順、提出時間の降順でソート
+    candidates.sort(key=lambda x: (x.get("_score", 0), x.get("submitDateTime", "")), reverse=True)
+    return candidates
+
+
+def download_edinet_pdf(doc_id: str) -> Optional[bytes]:
+    """
+    EDINET の docID から PDF バイト列を取得。
+    環境変数 EDINET_API_KEY が必要です。
+      - URL 文字列が渡された場合はそのまま GET します（簡易対応）。
+    """
+    if doc_id.startswith("http://") or doc_id.startswith("https://"):
+        r = _http_get(doc_id)
+        return r.content if r.status_code == 200 else None
+
+    api_key = os.environ.get("EDINET_API_KEY")
+    if not api_key:
+        # APIキー未設定の場合は UI 側でアップロードを利用してください
+        raise RuntimeError("EDINET_API_KEY is not set")
+
+    url = EDINET_API_DOC.format(doc_id=doc_id)
+    headers = {"X-API-KEY": api_key}
+    # type=1 が PDF
+    r = _http_get(url, headers=headers, params={"type": 1})
+    if r.status_code != 200:
+        return None
+    return r.content
+
+
+def ingest_edinet_for_company(edinet_code: str, date: str, max_docs: int = 2) -> int:
+    """
+    会社コード＋日付で EDINET doc を探し、PDFを取り込んでチャンク投入。
+    返値は投入したチャンク数。
+    """
+    if not edinet_code or not date:
+        raise ValueError("edinet_code と date は必須です（dateはYYYY-MM-DD）")
+
+    found = find_company_doc_ids(edinet_code, date, search_window_days=2)
+    if not found:
+        return 0
+
+    consumed = 0
+    for row in found[:max_docs]:
+        doc_id = row.get("docID") or row.get("docId")
+        if not doc_id:
+            continue
+        title = (row.get("title") or row.get("docDescription") or "edinet").strip()
+        pdf = download_edinet_pdf(doc_id)
+        if not pdf:
+            continue
+
+        text = pdf_to_text(pdf)
+        if not text.strip():
+            continue
+
+        recs = []
+        for j, chunk in enumerate(chunk_text(text)):
+            recs.append({
+                "text": chunk,
+                "source_url": f"/proxy/edinet/{doc_id}",
+                "title": title,
+                "doc_id": doc_id,
+                "chunk_id": f"{doc_id}-{j:04d}",
+            })
+        if recs:
+            add_to_index(recs)
+            consumed += len(recs)
+
+    return consumed
+
+
+# =============================================================================
+# アップロードPDFの取り込み
+# =============================================================================
+
+def ingest_pdf_bytes(title: str, source_url: str, pdf_bytes: bytes) -> int:
+    """
+    ローカル/アップロードPDFをインデックスへ投入。
+    """
     text = pdf_to_text(pdf_bytes)
     if not text.strip():
         return 0
@@ -17,3 +202,100 @@ def ingest_pdf_bytes(title: str, source_url: str, pdf_bytes: bytes):
     if recs:
         add_to_index(recs)
     return len(recs)
+
+
+# =============================================================================
+# PDF → テキスト抽出
+# =============================================================================
+
+def pdf_to_text(pdf_bytes: bytes) -> str:
+    """
+    PyMuPDF でプレーンテキスト抽出。日本語でも高精度。
+    画像だけのPDFはテキストが空になることがあります（OCRは未実装）。
+    """
+    try:
+        with fitz.open(stream=pdf_bytes, filetype="pdf") as doc:
+            texts = []
+            for page in doc:
+                # "text" は改行付きのレイアウトテキスト、"blocks"等でも可
+                t = page.get_text("text")
+                texts.append(t)
+        text = "\n\n".join(texts)
+        # 余計な全角空白の連続などを軽く正規化
+        text = re.sub(r"[ \t\u3000]+", " ", text)
+        # 連続改行の整理
+        text = re.sub(r"\n{3,}", "\n\n", text)
+        return text.strip()
+    except Exception:
+        return ""
+
+
+# =============================================================================
+# テキスト分割（日本語向けゆるふわチャンク）
+# =============================================================================
+
+_SENT_SPLIT_RE = re.compile(r"(.*?[\.\?\!。！？]\s*)", re.S)
+
+
+def split_sentences(text: str) -> List[str]:
+    """
+    「。！？.!?」までを1文として切り出し。末尾に句点が無い行も拾う。
+    """
+    parts = []
+    pos = 0
+    for m in _SENT_SPLIT_RE.finditer(text):
+        parts.append(m.group(0))
+        pos = m.end()
+    if pos < len(text):
+        parts.append(text[pos:])
+    # 行単位の段落も尊重（空行で分割したい場合はここで更に加工）
+    out = []
+    for p in parts:
+        p = p.strip()
+        if p:
+            out.extend([s for s in p.splitlines() if s.strip()])
+    return out
+
+
+def chunk_text(text: str, target_chars: int = 1000, overlap_chars: int = 200) -> List[str]:
+    """
+    文章境界をできるだけ保ちながら、概ね target_chars で分割。
+    チャンク間に overlap_chars の重なりをつけ、RAGの一致率を上げる。
+    """
+    if not text:
+        return []
+
+    sents = split_sentences(text)
+    chunks: List[str] = []
+    buf: List[str] = []
+    size = 0
+
+    def flush():
+        if not buf:
+            return
+        chunk = "".join(buf).strip()
+        if chunk:
+            chunks.append(chunk)
+
+    for s in sents:
+        if size + len(s) <= target_chars or not buf:
+            buf.append(s)
+            size += len(s)
+            continue
+        # いったん確定
+        flush()
+        # オーバーラップ確保
+        tail = []
+        remain = overlap_chars
+        # 後ろから文を足していく
+        for t in reversed(buf):
+            if remain <= 0:
+                break
+            tail.append(t)
+            remain -= len(t)
+        tail.reverse()
+        buf = tail + [s]
+        size = sum(len(x) for x in buf)
+
+    flush()
+    return chunks