Re-embed the knowledge base when the embedding model changes

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:
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
+52 -1
View File
@@ -9,6 +9,7 @@ from fastapi import APIRouter, Depends, HTTPException, Query
from pydantic import BaseModel
from sqlalchemy.orm import Session as DBSession
from app.core.memory_kb import local_knowledge
from app.core.memory_kb.local_knowledge import list_kb_files, read_kb_file
from app.db.database import get_db
from app.db.models import LocalKnowledgeDoc
@@ -41,10 +42,14 @@ class KbFileContent(BaseModel):
@router.get("/files", response_model=list[KbFileInfo])
def list_files(db: DBSession = Depends(get_db)):
"""ナレッジベース配下の Markdown を一覧で返す(本文は含めず、先頭のプレビューのみ)。"""
# 「埋め込み済み」は現在のモデルで埋め込まれている行だけを指す(別モデルの行は埋め直し待ち)
current_model = local_knowledge.current_embedding_model()
indexed_paths = {
# DB 側は OS 依存の区切り文字で保存されているため、比較用に posix 形式へ寄せる
path.replace("\\", "/")
for (path,) in db.query(LocalKnowledgeDoc.path).all()
for (path,) in db.query(LocalKnowledgeDoc.path)
.filter(LocalKnowledgeDoc.embedding_model == current_model)
.all()
}
return [
KbFileInfo(**info, indexed=info["path"] in indexed_paths)
@@ -59,3 +64,49 @@ def get_file(path: str = Query(..., description="ナレッジベースディレ
if file is None:
raise HTTPException(status_code=404, detail="Knowledge base file not found")
return KbFileContent(**file)
class EmbeddingStatus(BaseModel):
"""埋め込みモデルの現状。GUI の状態表示・進捗バーに使う。"""
model: str
# 現在のモデルで埋め込まれたベクトルの次元数(1件も無ければ None)
dimensions: int | None
file_count: int
indexed: int
# 別のモデルで埋め込まれていて埋め直しが必要な件数
stale: int
running: bool
total: int
done: int
failed: int
started_at: datetime | None
finished_at: datetime | None
last_error: str | None
class EmbeddingProbe(BaseModel):
"""短いテキストを実際に埋め込んでみた結果。"""
ok: bool
model: str
dimensions: int | None
latency_ms: float
error: str | None
@router.get("/embedding", response_model=EmbeddingStatus)
def get_embedding_status(db: DBSession = Depends(get_db)):
"""現在の埋め込みモデルと、埋め込み済み件数・実行中の進捗を返す。"""
return local_knowledge.embedding_status(db)
@router.post("/embedding/test", response_model=EmbeddingProbe)
async def test_embedding():
"""埋め込みモデルに実際に1回リクエストして、応答するかどうかを確かめる。"""
return await local_knowledge.probe_embedding()
@router.post("/embedding/reindex", response_model=EmbeddingStatus)
async def reindex_embeddings(db: DBSession = Depends(get_db)):
"""全ファイルの再埋め込みをバックグラウンドで開始する。実行中なら何もしない(進捗をそのまま返す)。"""
local_knowledge.start_reembed_all()
return local_knowledge.embedding_status(db)
+22 -2
View File
@@ -1,9 +1,12 @@
import os
from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel
from sqlalchemy.orm import Session as DBSession
from app.core import settings_store
from app.core.auth import require_settings_password
from app.core.memory_kb import local_knowledge
from app.db.database import get_db
router = APIRouter(prefix="/api/settings", tags=["settings"], dependencies=[Depends(require_settings_password)])
@@ -29,7 +32,8 @@ def list_settings(db: DBSession = Depends(get_db)):
@router.put("", response_model=list[SettingEntry])
def save_settings(req: SaveSettingsRequest, db: DBSession = Depends(get_db)):
async def save_settings(req: SaveSettingsRequest, db: DBSession = Depends(get_db)):
before_embedding_model = os.environ.get("LOCAL_KB_EMBEDDING_MODEL")
try:
for key, value in req.values.items():
if not value:
@@ -37,13 +41,29 @@ def save_settings(req: SaveSettingsRequest, db: DBSession = Depends(get_db)):
settings_store.save(db, key, value)
except ValueError as e:
raise HTTPException(status_code=400, detail=str(e)) from e
_reembed_if_model_changed(before_embedding_model)
return settings_store.list_masked(db)
@router.delete("/{key}", response_model=list[SettingEntry])
def clear_setting(key: str, db: DBSession = Depends(get_db)):
async def clear_setting(key: str, db: DBSession = Depends(get_db)):
before_embedding_model = os.environ.get("LOCAL_KB_EMBEDDING_MODEL")
try:
settings_store.clear(db, key)
except ValueError as e:
raise HTTPException(status_code=400, detail=str(e)) from e
_reembed_if_model_changed(before_embedding_model)
return settings_store.list_masked(db)
def _reembed_if_model_changed(before: str | None) -> None:
"""埋め込みモデルが変わったら、ローカルKBの再埋め込みをバックグラウンドで走らせる。
モデルが違うと次元数も意味空間も変わり、古いベクトルは比較しても無意味なので、
設定を変えた時点で全件やり直す(完了を待たずに応答は返す)。
"""
if os.environ.get("LOCAL_KB_EMBEDDING_MODEL") == before:
return
local_knowledge.start_reembed_all()
+1
View File
@@ -25,6 +25,7 @@ SESSION_MAX_AGE_DAYS: int = int(os.getenv("SESSION_MAX_AGE_DAYS", "30"))
TAGGING_MODEL: str | None = os.getenv("TAGGING_MODEL")
# Tavily Web検索
# GUI設定画面から実行時に変更できるため、実際の参照は search/websearch._api_key() が os.environ から行う。
TAVILY_API_KEY: str | None = os.getenv("TAVILY_API_KEY")
# 蓄積した知識をシステムプロンプトに注入する際の最大件数
KNOWLEDGE_CONTEXT_LIMIT: int = int(os.getenv("KNOWLEDGE_CONTEXT_LIMIT", "15"))
+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():
+14 -4
View File
@@ -1,13 +1,22 @@
import logging
import os
import httpx
from app.core.config import TAVILY_API_KEY
logger = logging.getLogger(__name__)
TAVILY_SEARCH_URL = "https://api.tavily.com/search"
def _api_key() -> str | None:
"""GUIの設定画面から実行時に変更できるよう、起動時定数ではなく都度 os.environ を見に行く。
app.core.config の定数は import 時に評価されるため、DBに保存した値を
settings_store.load_into_env() が os.environ に入れても反映されない。
"""
return os.environ.get("TAVILY_API_KEY") or None
WEB_SEARCH_TOOL = {
"type": "function",
"function": {
@@ -27,7 +36,8 @@ WEB_SEARCH_TOOL = {
async def search_web(query: str, max_results: int = 5) -> list[dict]:
"""Tavily APIでWeb検索する。結果は [{title, url, content}, ...]。失敗時は空リスト。"""
if not TAVILY_API_KEY:
api_key = _api_key()
if not api_key:
logger.warning("TAVILY_API_KEY is not set; skipping web search")
return []
@@ -36,7 +46,7 @@ async def search_web(query: str, max_results: int = 5) -> list[dict]:
response = await client.post(
TAVILY_SEARCH_URL,
json={
"api_key": TAVILY_API_KEY,
"api_key": api_key,
"query": query,
"max_results": max_results,
},
+1
View File
@@ -24,6 +24,7 @@ _ADDED_COLUMNS = [
("knowledge_entries", "headline", "VARCHAR"),
("messages", "archived_at", "DATETIME"),
("messages", "memory_id", "INTEGER"),
("local_knowledge_docs", "embedding_model", "VARCHAR"),
]
+5 -1
View File
@@ -104,7 +104,7 @@ class LocalKnowledgeDoc(Base):
"""backend/knowledge_base/ 配下のMarkdownファイルとその埋め込みベクトル。
話題に応じて関連ファイルをシステムプロンプトに注入するために使う。
ファイル内容のハッシュが変わった時だけ埋め込みを再計算する。
ファイル内容のハッシュか埋め込みモデルが変わった時だけ埋め込みを再計算する。
"""
__tablename__ = "local_knowledge_docs"
@@ -113,6 +113,10 @@ class LocalKnowledgeDoc(Base):
title: Mapped[str] = mapped_column(String)
content_hash: Mapped[str] = mapped_column(String)
embedding_json: Mapped[str] = mapped_column(Text)
# この埋め込みを計算したモデル名。モデルを切り替えた行を見分けて埋め直すために持つ
# (次元数の違うベクトル同士を比較しても意味のないスコアしか出ないため)。
# このカラムより前に作られた行は NULL で、モデル不明として埋め直し対象になる。
embedding_model: Mapped[str | None] = mapped_column(String, nullable=True)
updated_at: Mapped[datetime] = mapped_column(UtcDateTime, default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc))