Update irpr/deps.py

irpr/deps.py

@@ -1,157 +1,165 @@
-# irpr/deps.py  --- Chromadb版（faiss不使用・LLMなしでも動く）
+# irpr/deps.py  --- OpenAI埋め込み + 自前ベクタストア（numpy）／LLM生成
 from __future__ import annotations
-import os
-from typing import List, Dict, Optional
+import os, json, uuid
+from typing import List, Dict, Optional, Tuple
 import numpy as np
 from irpr.config import settings
 
-# ===== 書き込み可能な場所にキャッシュを集約 =====
-BASE = settings.DATA_DIR or "./var"
-DEFAULT_CACHE = os.path.join(BASE, ".hf-cache")
-DEFAULT_HOME  = os.path.join(BASE, ".hf-home")
+# ==== 書き込み可能ディレクトリの決定 ====
+def _pick_writable_dir() -> str:
+    candidates = [settings.DATA_DIR, "/data", "./var", "/tmp/irpr", "."]
+    for base in candidates:
+        try:
+            if not base: continue
+            os.makedirs(base, exist_ok=True)
+            p = os.path.join(base, ".write_test")
+            with open(p, "w") as w: w.write("ok")
+            os.remove(p)
+            return base
+        except Exception:
+            continue
+    return "."
+
+BASE_DIR = _pick_writable_dir()
+INDEX_DIR = settings.INDEX_DIR or os.path.join(BASE_DIR, "simple_index")
+os.makedirs(INDEX_DIR, exist_ok=True)
 
-# 既に設定が無ければ、ここで安全なデフォルトに設定
-os.environ.setdefault("HF_HOME", DEFAULT_HOME)
-os.environ.setdefault("TRANSFORMERS_CACHE", DEFAULT_CACHE)
-os.environ.setdefault("SENTENCE_TRANSFORMERS_HOME", DEFAULT_CACHE)
-os.environ.setdefault("HUGGINGFACE_HUB_CACHE", DEFAULT_CACHE)
+VECS_PATH = os.path.join(INDEX_DIR, "vectors.npy")     # np.float32 [N,D]（正規化済）
+META_PATH = os.path.join(INDEX_DIR, "meta.jsonl")       # 1行1メタ
+TEXT_PATH = os.path.join(INDEX_DIR, "texts.jsonl")      # 1行1テキスト
 
-for d in [os.environ["HF_HOME"], os.environ["TRANSFORMERS_CACHE"],
-          os.environ["SENTENCE_TRANSFORMERS_HOME"], os.environ["HUGGINGFACE_HUB_CACHE"],
-          BASE, settings.CHROMA_PATH]:
+# ==== OpenAI クライアント ====
+def _openai_client():
     try:
-        os.makedirs(d, exist_ok=True)
-    except Exception:
-        pass  # フォルダ作成に失敗しても致命ではない
+        from openai import OpenAI
+    except Exception as e:
+        raise RuntimeError("`openai` パッケージが見つかりません。requirements.txt に openai を追加してください。") from e
+    key = os.environ.get("OPENAI_API_KEY", "").strip()
+    if not key:
+        raise RuntimeError("OPENAI_API_KEY が未設定です。環境変数に設定してください。")
+    return OpenAI(api_key=key)
 
-# 遅延ロード
-_EMB = None
-_EMB_DIM: Optional[int] = None
-_CHROMA_COLLECTION = None
-_GEN = None
-_TOK = None
+# ==== 収納・ロード ====
+def _load_index() -> Tuple[np.ndarray, List[dict], List[str]]:
+    if os.path.exists(VECS_PATH):
+        vecs = np.load(VECS_PATH).astype(np.float32, copy=False)
+    else:
+        vecs = np.zeros((0, 0), dtype=np.float32)
+    metas: List[dict] = []
+    texts: List[str] = []
+    if os.path.exists(META_PATH):
+        with open(META_PATH, "r", encoding="utf-8") as f:
+            for line in f:
+                line = line.strip()
+                if line:
+                    metas.append(json.loads(line))
+    if os.path.exists(TEXT_PATH):
+        with open(TEXT_PATH, "r", encoding="utf-8") as f:
+            for line in f:
+                texts.append(line.rstrip("\n"))
+    # 整合性チェック
+    if vecs.size == 0:
+        return np.zeros((0, 0), dtype=np.float32), [], []
+    n = vecs.shape[0]
+    if len(metas) != n or len(texts) != n:
+        # 壊れているなら初期化
+        return np.zeros((0, 0), dtype=np.float32), [], []
+    return vecs, metas, texts
 
-def _get_embedder():
-    """SentenceTransformer を遅延ロード"""
-    global _EMB, _EMB_DIM
-    if _EMB is None:
-        from sentence_transformers import SentenceTransformer
-        _EMB = SentenceTransformer(
-            settings.EMB_MODEL,
-            cache_folder=os.environ.get("HF_HOME", DEFAULT_CACHE)
-        )
-        _EMB_DIM = _EMB.get_sentence_embedding_dimension()
-    return _EMB
+def _save_index(vecs: np.ndarray, metas: List[dict], texts: List[str]) -> None:
+    os.makedirs(INDEX_DIR, exist_ok=True)
+    np.save(VECS_PATH, vecs.astype(np.float32, copy=False))
+    with open(META_PATH, "w", encoding="utf-8") as f:
+        for m in metas:
+            f.write(json.dumps(m, ensure_ascii=False) + "\n")
+    with open(TEXT_PATH, "w", encoding="utf-8") as f:
+        for t in texts:
+            f.write((t or "").replace("\n", "\\n") + "\n")  # 1行1テキストに正規化
 
+# ==== Embedding ====
 def embed_texts(texts: List[str]) -> np.ndarray:
-    emb = _get_embedder()
-    arr = emb.encode(
-        texts,
-        batch_size=16,
-        normalize_embeddings=True,
-        convert_to_numpy=True,
-        show_progress_bar=False
-    )
-    return arr.astype(np.float32, copy=False)
-
-def _get_chroma():
-    """永続 Chromadb コレクションを取得"""
-    global _CHROMA_COLLECTION
-    if _CHROMA_COLLECTION is None:
-        import chromadb
-        from chromadb.config import Settings as CS
-        client = chromadb.PersistentClient(
-            path=settings.CHROMA_PATH,
-            settings=CS(allow_reset=True)
-        )
-        _CHROMA_COLLECTION = client.get_or_create_collection(name="irpr_docs")
-    return _CHROMA_COLLECTION
+    client = _openai_client()
+    model = settings.OPENAI_EMBED_MODEL
+    # バッチで呼ぶ
+    B = 128
+    out = []
+    for i in range(0, len(texts), B):
+        batch = texts[i:i+B]
+        resp = client.embeddings.create(model=model, input=batch)
+        out.extend([d.embedding for d in resp.data])
+    arr = np.array(out, dtype=np.float32)
+    # 正規化（コサイン類似度用）
+    norms = np.linalg.norm(arr, axis=1, keepdims=True) + 1e-12
+    return arr / norms
 
-def add_to_index(records: List[Dict]):
-    """records: {text, title, source_url, doc_id, chunk_id} の配列"""
+# ==== 追加 ====
+def add_to_index(records: List[Dict]) -> int:
+    """
+    records: [{text, title, source_url, doc_id, chunk_id}]
+    """
     if not records:
-        return
-    col = _get_chroma()
+        return 0
     texts = [r["text"] for r in records]
-    embs = embed_texts(texts)
+    vecs_new = embed_texts(texts)
+
+    vecs, metas, old_texts = _load_index()
+    if vecs.size == 0:
+        vecs = vecs_new
+        metas = []
+        old_texts = []
+    else:
+        if vecs.shape[1] != vecs_new.shape[1]:
+            # 埋め込み次元が違う（モデルを変えた等）→作り直し
+            vecs = vecs_new
+            metas = []
+            old_texts = []
+        else:
+            vecs = np.vstack([vecs, vecs_new])
 
-    ids, metas = [], []
     for r in records:
-        doc_id = r.get("doc_id") or "doc"
+        doc_id = r.get("doc_id") or str(uuid.uuid4())
         chunk_id = r.get("chunk_id") or ""
-        rid = f"{doc_id}:{chunk_id}" if chunk_id else doc_id
-        ids.append(rid)
         metas.append({
             "source_url": r.get("source_url"),
             "title": r.get("title"),
             "doc_id": doc_id,
             "chunk_id": chunk_id,
+            "id": f"{doc_id}:{chunk_id}" if chunk_id else doc_id
         })
-    col.add(ids=ids, documents=texts, embeddings=embs, metadatas=metas)
+        old_texts.append(r.get("text", ""))
+
+    _save_index(vecs, metas, old_texts)
+    return len(records)
 
+# ==== 検索 ====
 def search(query: str, top_k=8) -> List[Dict]:
-    col = _get_chroma()
-    q_emb = embed_texts([query])
-    res = col.query(
-        query_embeddings=q_emb,
-        n_results=top_k,
-        include=["documents", "metadatas", "distances", "ids"]
-    )
-    docs  = res.get("documents", [[]])[0]
-    metas = res.get("metadatas", [[]])[0]
-    dists = res.get("distances", [[]])[0]
+    vecs, metas, texts = _load_index()
+    if vecs.size == 0:
+        return []
+    q = embed_texts([query])[0]  # (D,)
+    scores = vecs @ q  # cosine (正規化済み)
+    idx = np.argsort(-scores)[:max(1, top_k)]
     out: List[Dict] = []
-    for text, meta, dist in zip(docs, metas, dists):
-        score = 1.0 - float(dist)/2.0 if dist is not None else None  # 類似度風スコア
+    for i in idx.tolist():
+        m = metas[i]
         out.append({
-            "text": text,
-            "source_url": (meta or {}).get("source_url"),
-            "title": (meta or {}).get("title"),
-            "doc_id": (meta or {}).get("doc_id"),
-            "chunk_id": (meta or {}).get("chunk_id"),
-            "score": score,
+            "text": (texts[i] or "").replace("\\n", "\n"),
+            "source_url": m.get("source_url"),
+            "title": m.get("title"),
+            "doc_id": m.get("doc_id"),
+            "chunk_id": m.get("chunk_id"),
+            "score": float(scores[i]),
         })
     return out
 
-# ==== 生成（任意） ====
-def _get_gen_pipeline():
-    """
-    GEN_MODEL が空なら LLM 無効の合図として例外を投げる。
-    CPU環境でも動くように dtype/device_map は保守的に。
-    """
-    if not settings.GEN_MODEL:
-        raise RuntimeError("GEN_MODEL is empty (LLM disabled).")
-    global _GEN, _TOK
-    if _GEN is None:
-        from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-        # torch は任意（無ければCPU既定）
-        try:
-            import torch  # noqa
-            torch_dtype = getattr(torch, "bfloat16", None)
-        except Exception:
-            torch = None
-            torch_dtype = None
-
-        name = settings.GEN_MODEL
-        cache_dir = os.environ.get("HF_HOME", DEFAULT_CACHE)
-        _TOK = AutoTokenizer.from_pretrained(name, cache_dir=cache_dir)
-        # dtype/device_map はCPUでも成立する保守的な指定にする
-        model_kwargs = dict(cache_dir=cache_dir, low_cpu_mem_usage=True)
-        if torch and hasattr(torch, "cuda") and torch.cuda.is_available():
-            model_kwargs["torch_dtype"] = getattr(torch, "bfloat16", None) or getattr(torch, "float16", None)
-            model_kwargs["device_map"] = "auto"
-
-        _MODEL = AutoModelForCausalLM.from_pretrained(name, **model_kwargs)
-        _GEN = pipeline("text-generation", model=_MODEL, tokenizer=_TOK)
-    return _GEN, _TOK
-
+# ==== 生成 ====
 def generate_chat(messages: List[Dict], max_new_tokens=600, temperature=0.2) -> str:
-    gen, tok = _get_gen_pipeline()
-    prompt = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    out = gen(
-        prompt,
-        do_sample=(temperature > 0.0),
-        temperature=temperature,
-        max_new_tokens=max_new_tokens
-    )[0]["generated_text"]
-    return out[len(prompt):].strip()
+    client = _openai_client()
+    model = settings.OPENAI_CHAT_MODEL
+    resp = client.chat.completions.create(
+        model=model,
+        messages=messages,
+        temperature=float(temperature),
+        max_tokens=int(max_new_tokens),
+    )
+    return (resp.choices[0].message.content or "").strip()