colcaroli/playchat
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Re-embed the knowledge base when the embedding model changes

Browse Source 
Vectors from two different embedding models are not comparable — Titan v2 has
1024 dimensions where nomic-embed-text has 768 — and _cosine_similarity zips
them, so a mismatch produced quietly meaningless scores rather than an error.
Nothing forced a re-embed either: sync_local_kb only looked at the content
hash, so switching models left every cached vector in place.

local_knowledge_docs now records the model each vector came from, and a row is
re-embedded when its content hash OR its model differs. Saving a new
LOCAL_KB_EMBEDDING_MODEL from the settings panel starts a background pass over
every file; the knowledge base panel also has a button for it. While that pass
runs, the per-request sync steps aside so no file is embedded twice.

The panel now shows what the embedding model is doing: model name, vector
size, how many files are embedded with the current model, how many are still
on an older one, a progress bar while a pass runs, and the last error. A
"接続テスト" button embeds one short string so a broken or unreachable model
can be told apart from one that is merely slow to warm up.

Also read TAVILY_API_KEY from os.environ at call time. It was imported as a
config constant, which is evaluated before load_into_env() runs, so a key
saved from the settings panel was stored but never used — not even after a
restart.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
colcaroli
2026-08-23 22:56:12 +09:00
co-authored by Claude Opus 5
parent 52dfc0c349
commit 97568ba90e
14 changed files with 550 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/app/core/memory_kb/local_knowledge.py
+178 -23
View File
@@ -3,11 +3,13 @@
- sync_local_kb(): ディレクトリをスキャンし、新規/変更されたファイルだけ埋め込みを計算してDBに保存する
- get_local_kb_context(): 直近のユーザー発言に埋め込み類似度の高いファイルを探し、システムプロンプト用テキストにする
"""
import asyncio
import hashlib
import json
import logging
import math
import os
import time
from datetime import datetime, timezone
from pathlib import Path
@@ -19,6 +21,46 @@ from app.db.models import KnowledgeEntry, LocalKnowledgeDoc
logger = logging.getLogger(__name__)
# 埋め込みの進行状況。GUIに出すためだけのプロセス内の状態で、永続化はしない
# （uvicornを複数ワーカーで動かすとワーカーごとに別の値になる点に注意）。
_progress: dict = {
"running": False,
"total": 0,
"done": 0,
"failed": 0,
"started_at": None,
"finished_at": None,
"last_error": None,
"last_model": None,
}
# create_task の戻り値。GCで途中終了しないように保持しておく
_reembed_task: "asyncio.Task | None" = None
def _start_progress(model: str, total: int) -> None:
_progress.update(
running=True, total=total, done=0, failed=0,
started_at=datetime.now(timezone.utc), finished_at=None,
last_error=None, last_model=model,
)
def _finish_progress() -> None:
_progress.update(running=False, finished_at=datetime.now(timezone.utc))
def _progress_snapshot() -> dict:
return {
"running": _progress["running"],
"total": _progress["total"],
"done": _progress["done"],
"failed": _progress["failed"],
"started_at": _progress["started_at"],
"finished_at": _progress["finished_at"],
"last_error": _progress["last_error"],
}
def _kb_dir() -> Path:
return Path(LOCAL_KB_DIR)
@@ -29,6 +71,11 @@ def _embedding_model() -> str:
return os.environ.get("LOCAL_KB_EMBEDDING_MODEL", "text-embedding-3-small")
def current_embedding_model() -> str:
"""いま使われている埋め込みモデル名。"""
return _embedding_model()
def _hash(content: str) -> str:
return hashlib.sha256(content.encode("utf-8")).hexdigest()
@@ -89,46 +136,75 @@ def backfill_search_results(db: DBSession) -> int:
return count
async def sync_local_kb(db: DBSession) -> None:
"""ナレッジベースディレクトリをスキャンし、新規/変更ファイルの埋め込みを更新する。
削除されたファイルに対応するレコードも削除する。
async def sync_local_kb(db: DBSession, force: bool = False) -> None:
"""ナレッジベースディレクトリをスキャンし、埋め込みが古いファイルだけ計算し直す。
埋め込みを取り直すのは、ファイル内容が変わったか、埋め込みモデルが変わったとき
（次元数の違うベクトルを比べても意味のあるスコアにならないため）。force=True なら全件やり直す。
削除されたファイルに対応するレコードも消す。
進捗は _progress に載せてGUIから見えるようにする。
"""
kb_dir = _kb_dir()
if not kb_dir.is_dir():
return
# 再埋め込みが走っている最中は、チャット毎の同期は譲って何もしない
# （同じファイルを二重に埋め込むのを避けるため。ここまでawaitを挟まないので取りこぼしはない）。
# force=True で呼ばれるのは start_reembed_all() が実行中フラグを立てた本人だけなので素通しする。
if _progress["running"] and not force:
return
model = _embedding_model()
md_files = sorted(kb_dir.rglob("*.md"))
seen_paths: set[str] = set()
# 先に対象を洗い出す。1件も無ければ進捗を触らない（毎リクエストの同期で履歴が消えないように）
pending: list[tuple[Path, str, str, LocalKnowledgeDoc | None]] = []
for file_path in md_files:
rel_path = str(file_path.relative_to(kb_dir))
seen_paths.add(rel_path)
content = file_path.read_text(encoding="utf-8")
content = file_path.read_text(encoding="utf-8", errors="replace")
content_hash = _hash(content)
existing = db.query(LocalKnowledgeDoc).filter(LocalKnowledgeDoc.path == rel_path).first()
if existing and existing.content_hash == content_hash:
up_to_date = (
existing is not None
and existing.content_hash == content_hash
and existing.embedding_model == model
)
if up_to_date and not force:
continue
pending.append((file_path, rel_path, content, existing))
if pending:
_start_progress(model, len(pending))
try:
[embedding] = await llm.embed([content], model=_embedding_model())
except Exception:
logger.exception("Failed to embed local knowledge file: %s", rel_path)
continue
for file_path, rel_path, content, existing in pending:
try:
[embedding] = await llm.embed([content], model=model)
except Exception as e:
logger.exception("Failed to embed local knowledge file: %s", rel_path)
_progress["failed"] += 1
_progress["last_error"] = f"{rel_path}: {e}"[:500]
continue
title = _title_from_content(content, fallback=file_path.stem)
if existing:
existing.content_hash = content_hash
existing.title = title
existing.embedding_json = json.dumps(embedding)
else:
db.add(LocalKnowledgeDoc(
path=rel_path,
title=title,
content_hash=content_hash,
embedding_json=json.dumps(embedding),
))
db.commit()
title = _title_from_content(content, fallback=file_path.stem)
if existing:
existing.content_hash = _hash(content)
existing.title = title
existing.embedding_json = json.dumps(embedding)
existing.embedding_model = model
else:
db.add(LocalKnowledgeDoc(
path=rel_path,
title=title,
content_hash=_hash(content),
embedding_json=json.dumps(embedding),
embedding_model=model,
))
db.commit()
_progress["done"] += 1
finally:
_finish_progress()
stale = db.query(LocalKnowledgeDoc).filter(LocalKnowledgeDoc.path.notin_(seen_paths)).all()
for doc in stale:
@@ -137,6 +213,85 @@ async def sync_local_kb(db: DBSession) -> None:
db.commit()
async def reembed_all() -> None:
"""全ファイルを現在のモデルで埋め込み直す。バックグラウンドタスクから呼ぶ想定で、DBセッションは自前で持つ。"""
from app.db.database import SessionLocal
db = SessionLocal()
try:
await sync_local_kb(db, force=True)
except Exception as e:
logger.exception("Re-embedding of the local knowledge base failed")
_progress["last_error"] = str(e)[:500]
finally:
# 対象0件で sync 側が進捗を触らなかった場合でも、実行中フラグは必ず下ろす
_finish_progress()
db.close()
def start_reembed_all() -> bool:
"""再埋め込みをバックグラウンドで開始する。すでに走っていれば False を返して何もしない。"""
global _reembed_task
if _progress["running"]:
return False
# GUIは開始リクエストの応答をそのまま状態表示に使うので、タスクが動き出すのを待たず
# ここで「実行中」にしておく（総数は実際にスキャンした時点で入る）。
_start_progress(_embedding_model(), 0)
# create_task の戻り値を保持しないとGCで途中終了することがあるためモジュールに持たせる
_reembed_task = asyncio.create_task(reembed_all())
return True
def embedding_status(db: DBSession) -> dict:
"""GUIの「埋め込み」表示用。現在のモデルと、そのモデルで埋め込み済みの件数を返す。"""
model = _embedding_model()
kb_dir = _kb_dir()
file_count = len(list(kb_dir.rglob("*.md"))) if kb_dir.is_dir() else 0
rows = db.query(LocalKnowledgeDoc.embedding_model, LocalKnowledgeDoc.embedding_json).all()
indexed = 0
dimensions = None
for row_model, embedding_json in rows:
if row_model != model:
continue
indexed += 1
if dimensions is None:
dimensions = len(json.loads(embedding_json))
return {
"model": model,
"dimensions": dimensions,
"file_count": file_count,
"indexed": indexed,
# 別のモデルで埋め込まれた（＝埋め直しが要る）行。モデル不明の古い行も含む
"stale": len(rows) - indexed,
**_progress_snapshot(),
}
async def probe_embedding() -> dict:
"""埋め込みモデルが実際に応答するかを短いテキストで確認する。GUIの「接続テスト」用。"""
model = _embedding_model()
start = time.perf_counter()
try:
[vector] = await llm.embed(["これは埋め込みモデルの疎通確認です。"], model=model)
except Exception as e:
return {
"ok": False,
"model": model,
"dimensions": None,
"latency_ms": round((time.perf_counter() - start) * 1000, 1),
"error": f"{type(e).__name__}: {e}"[:500],
}
return {
"ok": True,
"model": model,
"dimensions": len(vector),
"latency_ms": round((time.perf_counter() - start) * 1000, 1),
"error": None,
}
async def get_local_kb_context(db: DBSession, query_text: str) -> str | None:
"""直近のユーザー発言に関連するローカルMarkdown知識をシステムプロンプト用テキストにして返す。"""
if not query_text.strip():