diff --git a/src/server/app.py b/src/server/app.py index a4ade7d..5d4f6c7 100644 --- a/src/server/app.py +++ b/src/server/app.py @@ -21,14 +21,20 @@ logger = logging.getLogger("md-vector-db") class IngestRequest(BaseModel): file_path: str | None = None content: str | None = None - file_name: str | None = Field( - default=None, max_length=255 + file_name: str | None = Field(default=None, max_length=255) + collection: str | None = Field( + default=None, max_length=128, + description="目标 collection（默认使用配置文件中的 collection_name）", ) class SearchRequest(BaseModel): query: str = Field(..., min_length=1, max_length=2000) top_k: int = Field(default=10, ge=1, le=100) + collection: str | None = Field( + default=None, max_length=128, + description="检索的 collection（默认使用配置文件中的 collection_name）", + ) # -- 路径安全检查 -- @@ -64,9 +70,7 @@ def health(state: AppState = Depends(get_state)): @app.get("/api/v1/collections") def list_collections(state: AppState = Depends(get_state)): - info = state.searcher.get_collection_info() - sources = state.searcher.list_sources() - return {"collections": [info], "sources": sources} + return {"collections": state.list_collections_with_stats()} @app.post("/api/v1/ingest") @@ -75,6 +79,7 @@ def ingest_document( state: AppState = Depends(get_state), _: bool = Depends(verify_api_key), ): + ingestor = state.get_ingestor(req.collection) try: if req.file_path: if not _is_safe_path(req.file_path): @@ -85,14 +90,14 @@ def ingest_document( raise HTTPException(status_code=400, detail="不允许访问当前目录外的路径") if not path.exists(): raise HTTPException(status_code=404, detail=f"文件不存在: {path.name}") - count = state.ingestor.ingest_file(str(path)) + count = ingestor.ingest_file(str(path)) file_name = path.name elif req.content: file_name = req.file_name or "untitled.md" - count = state.ingestor.ingest_content(req.content, file_name) + count = ingestor.ingest_content(req.content, file_name) else: raise HTTPException(status_code=400, detail="需要提供 file_path 或 content") - return {"status": "ok", "chunks": count, "file": file_name} + return {"status": "ok", "chunks": count, "file": file_name, "collection": ingestor.collection_name} except HTTPException: raise except Exception: @@ -106,8 +111,9 @@ def search_documents( state: AppState = Depends(get_state), _: bool = Depends(verify_api_key), ): - results = state.searcher.search(req.query, top_k=req.top_k) - return {"results": results} + searcher = state.get_searcher(req.collection) + results = searcher.search(req.query, top_k=req.top_k) + return {"results": results, "collection": searcher.collection_name} @app.delete("/api/v1/documents/{file_name}") @@ -115,8 +121,10 @@ def delete_document( file_name: str, state: AppState = Depends(get_state), _: bool = Depends(verify_api_key), + collection: str | None = None, ): - deleted = state.searcher.delete_by_source(file_name) + searcher = state.get_searcher(collection) + deleted = searcher.delete_by_source(file_name) if not deleted: raise HTTPException(status_code=404, detail=f"文档不存在: {file_name}") - return {"status": "ok", "file": file_name} + return {"status": "ok", "file": file_name, "collection": searcher.collection_name} diff --git a/src/server/deps.py b/src/server/deps.py index 0e3746c..7525f60 100644 --- a/src/server/deps.py +++ b/src/server/deps.py @@ -12,7 +12,7 @@ logger = logging.getLogger(__name__) class AppState: - """应用级共享状态.""" + """应用级共享状态 — db 和 embedder 全局共享, searcher/ingestor 按 collection 懒创建.""" def __init__(self): config_path = os.environ.get("MD_VECTOR_CONFIG", "config.yaml") @@ -22,20 +22,48 @@ class AppState: "MD_VECTOR_DB_DATA_DIR", self.config.chroma.persist_dir ) self.db = VectorDB(persist_dir=data_dir) - self.embedder = create_embedder(self.config.embed) - collection = os.environ.get( + self.default_collection = os.environ.get( "MD_VECTOR_DB_COLLECTION", self.config.chroma.collection_name ) - self.searcher = Searcher(self.db, self.embedder, collection) - self.ingestor = DocumentIngestor(self.db, self.embedder, collection) + + # 按 collection 懒加载 searcher / ingestor + self._searchers: dict[str, Searcher] = {} + self._ingestors: dict[str, DocumentIngestor] = {} + + def get_searcher(self, collection: str | None = None) -> Searcher: + name = collection or self.default_collection + if name not in self._searchers: + self._searchers[name] = Searcher(self.db, self.embedder, name) + return self._searchers[name] + + def get_ingestor(self, collection: str | None = None) -> DocumentIngestor: + name = collection or self.default_collection + if name not in self._ingestors: + self._ingestors[name] = DocumentIngestor(self.db, self.embedder, name) + return self._ingestors[name] + + def list_collections_with_stats(self) -> list[dict]: + """列出所有 collection 及其统计.""" + result = [] + all_names = set(self._searchers.keys()) | set(self._ingestors.keys()) + all_names.add(self.default_collection) + for name in sorted(all_names): + try: + coll = self.db.get_or_create_collection(name) + result.append({"name": name, "count": coll.count()}) + except Exception: + result.append({"name": name, "count": 0}) + return result def is_healthy(self) -> dict: """真实健康检查: 验证 ChromaDB 和 Embedder 是否可用.""" status = {"status": "ok", "checks": {}} try: - count = self.searcher.collection.count() + count = self.db.get_or_create_collection( + self.default_collection + ).count() status["checks"]["chromadb"] = {"status": "ok", "count": count} except Exception as e: status["checks"]["chromadb"] = {"status": "error", "detail": str(e)} @@ -59,3 +87,8 @@ def get_state() -> AppState: logger.info("初始化应用状态...") _state = AppState() return _state + + +def get_default_collection() -> str: + """获取默认 collection 名, 供 CLI 等非 HTTP 场景使用.""" + return get_state().default_collection