"""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", ]