Update modules/rag_indexer.py

modules/rag_indexer.py

@@ -1,83 +1,114 @@
-import re, pickle
-from typing import List
-from pathlib import Path
+import os
+import json
 import requests
+from pathlib import Path
+from typing import List, Dict, Any, Optional
+
 from bs4 import BeautifulSoup
-from readability import Document
-import faiss
-from sentence_transformers import SentenceTransformer
-from modules.utils import ensure_dirs, chunk_text
 
-DATA_DIR = Path("data")
-INDEX_PATH = DATA_DIR / "vector_store.faiss"
-META_PATH = DATA_DIR / "vector_store_meta.pkl"
+# utils から書き込み先と分割関数を取得
+from .utils import ensure_dirs, data_dir, chunk_text
 
+# 依存は遅延ロード（モデル初期化は重いので）
 _model = None
 
 def _embedder():
+    """
+    SentenceTransformer を遅延初期化。
+    - すべてのキャッシュは utils.ensure_dirs() 側で /tmp 等の書き込み可パスへ固定済み
+    - 念のため cache_folder も明示して、/root 直下を一切使わない
+    """
     global _model
-    if _model is None:
-        _model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+    if _model is not None:
+        return _model
+
+    ensure_dirs()
+    cache_base = data_dir() / "hf_cache"
+
+    # ここで環境変数も最終確認（多重防御）
+    os.environ.setdefault("HF_HOME", str(data_dir() / "hf_home"))
+    os.environ.setdefault("HUGGINGFACE_HUB_CACHE", str(cache_base))
+    os.environ.setdefault("TRANSFORMERS_CACHE", str(cache_base))
+    os.environ.setdefault("SENTENCE_TRANSFORMERS_HOME", str(cache_base))
+    os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1")
+    os.environ.setdefault("HF_TOKEN", "")
+
+    from sentence_transformers import SentenceTransformer
+    model_name = os.getenv("EMBEDDING_MODEL", "sentence-transformers/all-MiniLM-L6-v2")
+    _model = SentenceTransformer(model_name, cache_folder=str(cache_base))
     return _model
 
-def _load_index():
-    if INDEX_PATH.exists():
-        index = faiss.read_index(str(INDEX_PATH))
-        with open(META_PATH, "rb") as f:
-            meta = pickle.load(f)
-        return index, meta
-    d = 384  # all-MiniLM-L6-v2
-    index = faiss.IndexFlatIP(d)
-    meta = []
-    return index, meta
-
-def _save_index(index, meta):
-    faiss.write_index(index, str(INDEX_PATH))
-    with open(META_PATH, "wb") as f:
-        pickle.dump(meta, f)
-
-def _extract_text_from_url(url: str) -> str:
+def _write_chunks(rows: List[Dict[str, Any]]) -> int:
+    """
+    chunks.jsonl に追記（既存は破棄せず、簡易に追加）
+    1行: {"text": "...", "source": "path_or_url", "meta": {...}}
+    """
+    ensure_dirs()
+    out = data_dir() / "chunks.jsonl"
+    with open(out, "a", encoding="utf-8") as f:
+        for r in rows:
+            f.write(json.dumps(r, ensure_ascii=False) + "\n")
+    return len(rows)
+
+def _load_text_from_url(url: str) -> str:
     try:
-        r = requests.get(url, timeout=20, headers={"User-Agent":"Mozilla/5.0"})
+        r = requests.get(url, timeout=15)
         r.raise_for_status()
-        doc = Document(r.text)
-        html = doc.summary()
-        soup = BeautifulSoup(html, "lxml")
-        text = soup.get_text("\n")
-        return re.sub(r"\n{2,}", "\n", text).strip()
-    except Exception as e:
-        return f"[ERROR] failed to fetch {url}: {e}"
+        html = r.text
+        soup = BeautifulSoup(html, "html.parser")
+        # タイトル + 本文テキスト（簡易）
+        title = (soup.title.string.strip() if soup.title and soup.title.string else "")
+        text = soup.get_text("\n", strip=True)
+        return (title + "\n\n" + text).strip()
+    except Exception:
+        return ""
 
-def _extract_text_from_file(path: str) -> str:
-    p = Path(path)
-    if not p.exists():
+def _load_text_from_file(path: Path) -> str:
+    # テキスト/Markdown想定（PDF等は最小構成では未対応）
+    try:
+        with open(path, "r", encoding="utf-8", errors="ignore") as f:
+            return f.read()
+    except Exception:
         return ""
-    if p.suffix.lower() in [".txt", ".md", ".csv", ".json", ".py"]:
-        return p.read_text(errors="ignore")
-    return f"[FILE]{p.name}"
 
-def index_files_and_urls(file_paths: List[str], urls: List[str]) -> str:
+def index_files_and_urls(file_paths: Optional[List[str]] = None, urls: Optional[List[str]] = None) -> str:
+    """
+    - 受け取ったファイルとURLからテキストを抽出し、チャンク化して chunks.jsonl に追記
+    - 依存を最小化（PDF/Officeは最小構成では対象外）
+    """
     ensure_dirs()
-    index, meta = _load_index()
-    emb = _embedder()
-
-    docs = []
-    for u in urls or []:
-        text = _extract_text_from_url(u)
-        if text:
-            docs.append((u, text))
-    for fp in file_paths or []:
-        text = _extract_text_from_file(fp)
-        if text:
-            docs.append((fp, text))
+    file_paths = file_paths or []
+    urls = urls or []
 
     added = 0
-    for src, text in docs:
-        for chunk in chunk_text(text, 600):
-            vec = emb.encode([chunk], normalize_embeddings=True)
-            index.add(vec)
-            meta.append({"source": src, "text": chunk})
-            added += 1
-
-    _save_index(index, meta)
-    return f"Indexed {added} chunks from {len(docs)} sources."
+    rows: List[Dict[str, Any]] = []
+
+    # ファイル
+    for p in file_paths:
+        try:
+            path = Path(p)
+            txt = _load_text_from_file(path)
+            for ch in chunk_text(txt):
+                rows.append({"text": ch, "source": str(path), "meta": {"kind": "file"}})
+        except Exception:
+            continue
+
+    # URL
+    for u in urls:
+        txt = _load_text_from_url(u)
+        for ch in chunk_text(txt):
+            rows.append({"text": ch, "source": u, "meta": {"kind": "url"}})
+
+    if rows:
+        added = _write_chunks(rows)
+
+    # 埋め込みモデルを一度初期化しておく（初回ダウンロードが必要なため）
+    try:
+        emb = _embedder()
+        # 1件だけ実行してキャッシュを準備
+        _ = emb.encode(["warmup"], normalize_embeddings=True)
+        warmed = True
+    except Exception as e:
+        warmed = False
+
+    return f"indexed_chunks={added}, warmed_up={warmed}"