107 lines
3.9 KiB
Python
107 lines
3.9 KiB
Python
"""FastAPI 依赖注入 — 集中管理应用状态, 替代模块级全局变量."""
|
|
import logging
|
|
import os
|
|
import threading
|
|
|
|
from src.core.config import load_config
|
|
from src.core.db import VectorDB
|
|
from src.core.embedder import create_embedder
|
|
from src.core.ingest import DocumentIngestor
|
|
from src.core.search import Searcher
|
|
|
|
logger = logging.getLogger("md-vector-db")
|
|
|
|
|
|
class AppState:
|
|
"""应用级共享状态 — db 和 embedder 全局共享, searcher/ingestor 按 collection 懒创建."""
|
|
|
|
def __init__(self):
|
|
config_path = os.environ.get("MD_VECTOR_CONFIG", "config.yaml")
|
|
self.config = load_config(config_path)
|
|
|
|
data_dir = os.environ.get(
|
|
"MD_VECTOR_DB_DATA_DIR", self.config.chroma.persist_dir
|
|
)
|
|
self.db = VectorDB(persist_dir=data_dir)
|
|
self.embedder = create_embedder(self.config.embed)
|
|
|
|
self.default_collection = os.environ.get(
|
|
"MD_VECTOR_DB_COLLECTION", self.config.chroma.collection_name
|
|
)
|
|
|
|
# 按 collection 懒加载 searcher / ingestor
|
|
self._searchers: dict[str, Searcher] = {}
|
|
self._ingestors: dict[str, DocumentIngestor] = {}
|
|
self._cache_lock = threading.Lock()
|
|
|
|
def get_searcher(self, collection: str | None = None) -> Searcher:
|
|
name = collection or self.default_collection
|
|
with self._cache_lock:
|
|
if name not in self._searchers:
|
|
self._searchers[name] = Searcher(
|
|
self.db, self.embedder, name,
|
|
search_config=self.config.search,
|
|
)
|
|
return self._searchers[name]
|
|
|
|
def get_ingestor(self, collection: str | None = None) -> DocumentIngestor:
|
|
name = collection or self.default_collection
|
|
with self._cache_lock:
|
|
if name not in self._ingestors:
|
|
self._ingestors[name] = DocumentIngestor(
|
|
self.db, self.embedder, name,
|
|
chunk_config=self.config.chunk,
|
|
)
|
|
return self._ingestors[name]
|
|
|
|
def list_collections_with_stats(self) -> list[dict]:
|
|
"""列出所有 collection 及其统计（直接从 ChromaDB 查询）."""
|
|
result = []
|
|
for coll in self.db.list_collections():
|
|
result.append({"name": coll.name, "count": coll.count()})
|
|
return result
|
|
|
|
def is_healthy(self) -> dict:
|
|
"""真实健康检查: 验证 ChromaDB 和 Embedder 是否可用.
|
|
|
|
对外仅返回 ok/degraded 状态，详细信息记入日志，不泄露内部路径.
|
|
"""
|
|
status = {"status": "ok", "checks": {}}
|
|
try:
|
|
count = self.db.get_or_create_collection(
|
|
self.default_collection
|
|
).count()
|
|
status["checks"]["chromadb"] = {"status": "ok", "count": count}
|
|
except Exception as e:
|
|
logger.error("ChromaDB 健康检查失败: %s", e)
|
|
status["checks"]["chromadb"] = {"status": "error", "detail": "unavailable"}
|
|
status["status"] = "degraded"
|
|
try:
|
|
dim = self.embedder.dimension
|
|
status["checks"]["embedder"] = {"status": "ok", "dimension": dim}
|
|
except Exception as e:
|
|
logger.error("Embedder 健康检查失败: %s", e)
|
|
status["checks"]["embedder"] = {"status": "error", "detail": "unavailable"}
|
|
status["status"] = "degraded"
|
|
return status
|
|
|
|
|
|
_state: AppState | None = None
|
|
_state_lock = threading.Lock()
|
|
|
|
|
|
def get_state() -> AppState:
|
|
"""获取应用状态单例 (懒初始化, 线程安全)."""
|
|
global _state
|
|
if _state is None:
|
|
with _state_lock:
|
|
if _state is None:
|
|
logger.info("初始化应用状态...")
|
|
_state = AppState()
|
|
return _state
|
|
|
|
|
|
def get_default_collection() -> str:
|
|
"""获取默认 collection 名, 供 CLI 等非 HTTP 场景使用."""
|
|
return get_state().default_collection
|