372 lines
14 KiB
Python
372 lines
14 KiB
Python
"""Hot-swappable runtime composition for the DeerFlow-local Wiki index.
|
|
|
|
The file configuration remains the source of truth for the normal embedding
|
|
provider. Administrators can temporarily select the uploaded offline BGE-M3
|
|
encoder through ``system_settings.json`` without rewriting release assets or
|
|
restarting Gateway.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from fastapi import FastAPI
|
|
|
|
from deerflow.config import get_app_config
|
|
from deerflow.config.app_config import AppConfig
|
|
from deerflow.config.llmwiki_config import LocalWikiIndexConfig
|
|
from deerflow.config.system_settings import LocalWikiEncoderSettings, load_system_settings
|
|
from deerflow.integrations.weknora.runtime import build_weknora_client, get_resolved_llmwiki_runtime
|
|
|
|
logger = logging.getLogger(__name__)
|
|
_EMBEDDING_RETIREMENT_GRACE_SECONDS = 600.0
|
|
|
|
|
|
def resolve_effective_local_wiki_index_config(
|
|
config: AppConfig,
|
|
encoder_settings: LocalWikiEncoderSettings,
|
|
) -> LocalWikiIndexConfig:
|
|
"""Apply the persisted encoder mode to a copy of the file configuration."""
|
|
|
|
local_index = config.llmwiki.local_wiki_index
|
|
configured_embedding = local_index.embedding
|
|
runtime_model_path = encoder_settings.model_path.strip()
|
|
|
|
if encoder_settings.mode == "local":
|
|
# Directory upload is intentionally a BGE-M3 contract. A configured
|
|
# OpenAI model/dimension must not leak into the local FastEmbed worker.
|
|
embedding = configured_embedding.model_copy(
|
|
update={
|
|
"enabled": True,
|
|
"provider": "local",
|
|
"model": "bge-embedding-m3",
|
|
"dimensions": 1024,
|
|
"base_url": "",
|
|
"api_key": "",
|
|
"local_model_path": runtime_model_path or configured_embedding.local_model_path,
|
|
"local_model_revision": encoder_settings.model_revision,
|
|
}
|
|
)
|
|
elif configured_embedding.provider == "local":
|
|
# Keep the historical invariant: an administrator-selected directory
|
|
# overrides config.yaml when the configured provider itself is local.
|
|
updates = {"local_model_revision": encoder_settings.model_revision}
|
|
if runtime_model_path:
|
|
updates["local_model_path"] = runtime_model_path
|
|
embedding = configured_embedding.model_copy(
|
|
update=updates
|
|
)
|
|
else:
|
|
embedding = configured_embedding
|
|
return local_index.model_copy(update={"embedding": embedding})
|
|
|
|
|
|
def _resolved_local_model_path(config: LocalWikiIndexConfig) -> Path | None:
|
|
raw = config.embedding.local_model_path.strip() or os.environ.get(
|
|
"DEERFLOW_BGE_M3_MODEL_PATH", ""
|
|
).strip()
|
|
if not raw:
|
|
return None
|
|
path = Path(raw).expanduser()
|
|
return path.resolve() if path.is_absolute() else (Path.cwd() / path).resolve()
|
|
|
|
|
|
def _local_model_files_ready(config: LocalWikiIndexConfig) -> bool:
|
|
path = _resolved_local_model_path(config)
|
|
return bool(
|
|
path
|
|
and path.is_dir()
|
|
and ((path / "onnx" / "model.onnx").is_file() or (path / "model.onnx").is_file())
|
|
)
|
|
|
|
|
|
def _embedding_provider_ready(config: LocalWikiIndexConfig) -> bool:
|
|
embedding = config.embedding
|
|
if embedding.provider == "local":
|
|
return _local_model_files_ready(config)
|
|
return bool(embedding.base_url)
|
|
|
|
|
|
def _schedule_embedding_retirement(app: FastAPI, embedding: Any) -> None:
|
|
retired = list(getattr(app.state, "llmwiki_retired_embeddings", []) or [])
|
|
if embedding not in retired:
|
|
retired.append(embedding)
|
|
app.state.llmwiki_retired_embeddings = retired
|
|
|
|
async def retire_after_grace() -> None:
|
|
retired_successfully = False
|
|
try:
|
|
await asyncio.sleep(_EMBEDDING_RETIREMENT_GRACE_SECONDS)
|
|
await embedding.retire()
|
|
retired_successfully = True
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception:
|
|
logger.exception("Failed to retire an old Local Wiki embedding runtime")
|
|
finally:
|
|
if retired_successfully:
|
|
current = list(getattr(app.state, "llmwiki_retired_embeddings", []) or [])
|
|
app.state.llmwiki_retired_embeddings = [item for item in current if item is not embedding]
|
|
|
|
tasks: set[asyncio.Task[None]] = getattr(app.state, "llmwiki_retirement_tasks", None)
|
|
if tasks is None:
|
|
tasks = set()
|
|
app.state.llmwiki_retirement_tasks = tasks
|
|
task = asyncio.create_task(retire_after_grace(), name="llmwiki-embedding-retirement")
|
|
tasks.add(task)
|
|
task.add_done_callback(tasks.discard)
|
|
|
|
|
|
def _runtime_signature(
|
|
local_index: LocalWikiIndexConfig,
|
|
*,
|
|
resolved_llmwiki: Any,
|
|
encoder_mode: str,
|
|
) -> str:
|
|
payload = {
|
|
"local_index": local_index.model_dump(mode="json"),
|
|
# A directory can be mounted/unmounted without editing config.yaml.
|
|
# Readiness is part of the generation so services appear/disappear on
|
|
# the next lightweight scheduler poll instead of requiring a restart.
|
|
"local_model_files_ready": (
|
|
_local_model_files_ready(local_index)
|
|
if local_index.embedding.provider == "local"
|
|
else None
|
|
),
|
|
"encoder_mode": encoder_mode,
|
|
"weknora": {
|
|
"provider": resolved_llmwiki.provider,
|
|
"api_base_url": resolved_llmwiki.api_base_url,
|
|
"web_base_url": resolved_llmwiki.web_base_url,
|
|
"admin_email": resolved_llmwiki.admin_email,
|
|
"admin_password": resolved_llmwiki.admin_password,
|
|
"override_enabled": resolved_llmwiki.override_enabled,
|
|
},
|
|
}
|
|
raw = json.dumps(payload, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
|
|
return hashlib.sha256(raw.encode("utf-8")).hexdigest()
|
|
|
|
|
|
async def rebuild_local_wiki_index_runtime(
|
|
app: FastAPI,
|
|
*,
|
|
force: bool = False,
|
|
) -> LocalWikiIndexConfig:
|
|
"""Refresh config and atomically replace all embedding-dependent services."""
|
|
|
|
lock = getattr(app.state, "llmwiki_runtime_lock", None)
|
|
if lock is None:
|
|
lock = asyncio.Lock()
|
|
app.state.llmwiki_runtime_lock = lock
|
|
|
|
async with lock:
|
|
source_config = get_app_config()
|
|
settings = load_system_settings()
|
|
effective_config = resolve_effective_local_wiki_index_config(
|
|
source_config,
|
|
settings.local_wiki_encoder,
|
|
)
|
|
resolved_llmwiki = get_resolved_llmwiki_runtime(source_config)
|
|
signature = _runtime_signature(
|
|
effective_config,
|
|
resolved_llmwiki=resolved_llmwiki,
|
|
encoder_mode=settings.local_wiki_encoder.mode,
|
|
)
|
|
current_signature = getattr(app.state, "llmwiki_runtime_signature", None)
|
|
if not force and current_signature == signature:
|
|
app.state.llmwiki_source_config = source_config
|
|
app.state.llmwiki_effective_config = effective_config
|
|
return effective_config
|
|
|
|
index_store = getattr(app.state, "llmwiki_index_store", None)
|
|
if index_store is None:
|
|
raise RuntimeError("Local Wiki index store is not initialized")
|
|
|
|
from deerflow.integrations.weknora.local_index.embedding import StrictWikiEmbeddingClient
|
|
from deerflow.integrations.weknora.local_index.retrieval import WikiRetrievalService
|
|
from deerflow.integrations.weknora.local_index.runtime import register_local_wiki_index_runtime
|
|
from deerflow.integrations.weknora.local_index.search import WikiVectorSearchService
|
|
from deerflow.integrations.weknora.local_index.sync import WikiSyncService
|
|
from deerflow.integrations.weknora.local_index.vector_cache import WikiVectorCache
|
|
|
|
embedding_config = effective_config.embedding
|
|
provider_ready = _embedding_provider_ready(effective_config)
|
|
embedding_configured = bool(
|
|
embedding_config.enabled
|
|
and provider_ready
|
|
and embedding_config.model.strip()
|
|
)
|
|
old_embedding = getattr(app.state, "llmwiki_embedding", None)
|
|
reusable = (
|
|
old_embedding
|
|
if old_embedding is not None and old_embedding.config == embedding_config
|
|
else None
|
|
)
|
|
new_embedding = None
|
|
if embedding_configured:
|
|
new_embedding = reusable or StrictWikiEmbeddingClient(embedding_config)
|
|
new_cache = None
|
|
new_search = None
|
|
new_sync = None
|
|
new_retrieval = None
|
|
weknora_client = (
|
|
build_weknora_client(resolved_llmwiki)
|
|
if resolved_llmwiki.weknora_enabled
|
|
else None
|
|
)
|
|
|
|
if effective_config.enabled and new_embedding is not None:
|
|
new_cache = WikiVectorCache(
|
|
index_store,
|
|
max_knowledge_bases=effective_config.cache_max_knowledge_bases,
|
|
revision_check_seconds=effective_config.cache_revision_check_seconds,
|
|
)
|
|
new_search = WikiVectorSearchService(
|
|
effective_config,
|
|
new_embedding,
|
|
new_cache,
|
|
)
|
|
if weknora_client is not None:
|
|
new_sync = WikiSyncService(
|
|
effective_config,
|
|
index_store,
|
|
new_embedding,
|
|
new_cache,
|
|
weknora_client,
|
|
)
|
|
from app.gateway.knowledge_transfer import sync_to_global
|
|
|
|
async def on_wiki_vectorized(mapping):
|
|
return await sync_to_global(app, mapping)
|
|
|
|
new_sync.on_completed = on_wiki_vectorized
|
|
|
|
if weknora_client is not None or new_search is not None:
|
|
new_retrieval = WikiRetrievalService(
|
|
index_store,
|
|
vector_search=(new_search if effective_config.internal_search_enabled else None),
|
|
client=weknora_client,
|
|
title_boost_enabled=effective_config.title_boost_enabled,
|
|
title_boost_weight=effective_config.title_boost_weight,
|
|
)
|
|
|
|
app.state.llmwiki_source_config = source_config
|
|
app.state.llmwiki_effective_config = effective_config
|
|
app.state.llmwiki_embedding = new_embedding
|
|
app.state.llmwiki_vector_cache = new_cache
|
|
app.state.llmwiki_vector_search = new_search
|
|
app.state.llmwiki_sync_service = new_sync
|
|
app.state.llmwiki_retrieval_service = new_retrieval
|
|
app.state.llmwiki_runtime_signature = signature
|
|
register_local_wiki_index_runtime(
|
|
store=index_store,
|
|
search=new_search,
|
|
sync=new_sync,
|
|
retrieval=new_retrieval,
|
|
)
|
|
|
|
if old_embedding is not None and old_embedding is not new_embedding:
|
|
# Give requests holding the prior service graph time to finish,
|
|
# then reject any stragglers and release the potentially multi-GB
|
|
# local worker instead of accumulating one process per upload.
|
|
_schedule_embedding_retirement(app, old_embedding)
|
|
return effective_config
|
|
|
|
|
|
async def close_local_wiki_index_runtime(app: FastAPI) -> None:
|
|
"""Close the current and safely-retired embedding clients at shutdown."""
|
|
|
|
retired = list(getattr(app.state, "llmwiki_retired_embeddings", []) or [])
|
|
retirement_tasks = list(getattr(app.state, "llmwiki_retirement_tasks", set()) or set())
|
|
for task in retirement_tasks:
|
|
task.cancel()
|
|
if retirement_tasks:
|
|
await asyncio.gather(*retirement_tasks, return_exceptions=True)
|
|
app.state.llmwiki_retirement_tasks = set()
|
|
current = getattr(app.state, "llmwiki_embedding", None)
|
|
seen: set[int] = set()
|
|
for embedding in [current, *retired]:
|
|
if embedding is None or id(embedding) in seen:
|
|
continue
|
|
seen.add(id(embedding))
|
|
await embedding.close()
|
|
app.state.llmwiki_retired_embeddings = []
|
|
|
|
|
|
def local_encoder_runtime_status(app: FastAPI) -> dict[str, object]:
|
|
"""Return one stable status contract for remote and local providers."""
|
|
|
|
source_config: AppConfig | None = getattr(app.state, "llmwiki_source_config", None)
|
|
if source_config is None:
|
|
source_config = app.state.config
|
|
configured = source_config.llmwiki.local_wiki_index.embedding
|
|
effective: LocalWikiIndexConfig = getattr(
|
|
app.state,
|
|
"llmwiki_effective_config",
|
|
source_config.llmwiki.local_wiki_index,
|
|
)
|
|
settings = load_system_settings().local_wiki_encoder
|
|
embedding = getattr(app.state, "llmwiki_embedding", None)
|
|
if embedding is not None:
|
|
status = dict(embedding.local_encoder_status())
|
|
else:
|
|
status = {
|
|
"provider": effective.embedding.provider,
|
|
"model": effective.embedding.model,
|
|
"dimensions": effective.embedding.dimensions,
|
|
"state": "idle",
|
|
"running": False,
|
|
"last_error": None,
|
|
"started_at": None,
|
|
"ready_at": None,
|
|
"offline_only": effective.embedding.provider == "local",
|
|
"model_path_configured": False,
|
|
"model_path_exists": False,
|
|
"model_files_ready": False,
|
|
"process_isolation": effective.embedding.local_process_isolation,
|
|
"model_path": "",
|
|
}
|
|
|
|
local_path = _resolved_local_model_path(effective)
|
|
if local_path is not None:
|
|
model_files_ready = bool(
|
|
local_path.is_dir()
|
|
and (
|
|
(local_path / "onnx" / "model.onnx").is_file()
|
|
or (local_path / "model.onnx").is_file()
|
|
)
|
|
)
|
|
status.update(
|
|
{
|
|
"model_path": str(local_path),
|
|
"model_path_configured": True,
|
|
"model_path_exists": local_path.is_dir(),
|
|
"model_files_ready": model_files_ready,
|
|
}
|
|
)
|
|
|
|
status.update(
|
|
{
|
|
"provider": effective.embedding.provider,
|
|
"configured_provider": configured.provider,
|
|
"mode": settings.mode,
|
|
"using_local": effective.embedding.provider == "local",
|
|
"offline_only": effective.embedding.provider == "local",
|
|
}
|
|
)
|
|
return status
|
|
|
|
|
|
__all__ = [
|
|
"close_local_wiki_index_runtime",
|
|
"local_encoder_runtime_status",
|
|
"rebuild_local_wiki_index_runtime",
|
|
"resolve_effective_local_wiki_index_config",
|
|
]
|