Update rag/ingest.py

rag/ingest.py

@@ -1,301 +1,45 @@
 # rag/ingest.py
-"""
-EDINET/アップロードPDFの取り込み（RAGインデックス投入）ユーティリティ
-
-- 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 io, uuid
+from typing import List
 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
+def _split_text(text: str, chunk_size=800, overlap=150) -> List[str]:
+    text = (text or "").strip()
+    if not text:
+        return []
+    chunks = []
+    i = 0
+    while i < len(text):
+        chunk = text[i:i+chunk_size]
+        chunks.append(chunk)
+        i += chunk_size - overlap
+        if i < 0 or i >= len(text):
+            break
+    return chunks
 
-        recs = []
-        for j, chunk in enumerate(chunk_text(text)):
-            recs.append({
-                "text": chunk,
-                "source_url": f"/proxy/edinet/{doc_id}",
+def ingest_pdf_bytes(title: str, source_url: str, pdf_bytes: bytes) -> int:
+    doc = fitz.open(stream=io.BytesIO(pdf_bytes), filetype="pdf")
+    all_chunks = []
+    doc_id = str(uuid.uuid4())
+    for page_no in range(doc.page_count):
+        page = doc.load_page(page_no)
+        raw = page.get_text("text")
+        # ページ番号などを付与しておく
+        page_text = f"[p.{page_no+1}] {raw}".strip()
+        for j, ch in enumerate(_split_text(page_text, 900, 150)):
+            all_chunks.append({
+                "text": ch,
                 "title": title,
+                "source_url": source_url,
                 "doc_id": doc_id,
-                "chunk_id": f"{doc_id}-{j:04d}",
+                "chunk_id": f"{page_no+1}-{j+1}",
             })
-        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
-    recs = []
-    for j, chunk in enumerate(chunk_text(text)):
-        recs.append({
-            "text": chunk,
-            "source_url": source_url or "upload",
-            "title": title or "upload",
-            "doc_id": f"upload:{title}",
-            "chunk_id": f"upload-{j:04d}"
-        })
-    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 ""
-
+    doc.close()
+    return add_to_index(all_chunks)
 
-# =============================================================================
-# テキスト分割（日本語向けゆるふわチャンク）
-# =============================================================================
-
-_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
+# ---- EDINET ダミー実装（OpenAI専用版では未サポート）----
+def ingest_edinet_for_company(edinet_code: str, date: str) -> int:
+    # ここでは何もしない（将来実装用の置き場所）
+    return 0