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:
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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"))
|
||||
|
||||
@@ -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():
|
||||
|
||||
@@ -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,
|
||||
},
|
||||
|
||||
@@ -24,6 +24,7 @@ _ADDED_COLUMNS = [
|
||||
("knowledge_entries", "headline", "VARCHAR"),
|
||||
("messages", "archived_at", "DATETIME"),
|
||||
("messages", "memory_id", "INTEGER"),
|
||||
("local_knowledge_docs", "embedding_model", "VARCHAR"),
|
||||
]
|
||||
|
||||
|
||||
|
||||
@@ -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))
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user