feat: multi-collection support — each project uses its own isolated collection
This commit is contained in:
+20
-12
@@ -21,14 +21,20 @@ logger = logging.getLogger("md-vector-db")
|
|||||||
class IngestRequest(BaseModel):
|
class IngestRequest(BaseModel):
|
||||||
file_path: str | None = None
|
file_path: str | None = None
|
||||||
content: str | None = None
|
content: str | None = None
|
||||||
file_name: str | None = Field(
|
file_name: str | None = Field(default=None, max_length=255)
|
||||||
default=None, max_length=255
|
collection: str | None = Field(
|
||||||
|
default=None, max_length=128,
|
||||||
|
description="目标 collection(默认使用配置文件中的 collection_name)",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
class SearchRequest(BaseModel):
|
class SearchRequest(BaseModel):
|
||||||
query: str = Field(..., min_length=1, max_length=2000)
|
query: str = Field(..., min_length=1, max_length=2000)
|
||||||
top_k: int = Field(default=10, ge=1, le=100)
|
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")
|
@app.get("/api/v1/collections")
|
||||||
def list_collections(state: AppState = Depends(get_state)):
|
def list_collections(state: AppState = Depends(get_state)):
|
||||||
info = state.searcher.get_collection_info()
|
return {"collections": state.list_collections_with_stats()}
|
||||||
sources = state.searcher.list_sources()
|
|
||||||
return {"collections": [info], "sources": sources}
|
|
||||||
|
|
||||||
|
|
||||||
@app.post("/api/v1/ingest")
|
@app.post("/api/v1/ingest")
|
||||||
@@ -75,6 +79,7 @@ def ingest_document(
|
|||||||
state: AppState = Depends(get_state),
|
state: AppState = Depends(get_state),
|
||||||
_: bool = Depends(verify_api_key),
|
_: bool = Depends(verify_api_key),
|
||||||
):
|
):
|
||||||
|
ingestor = state.get_ingestor(req.collection)
|
||||||
try:
|
try:
|
||||||
if req.file_path:
|
if req.file_path:
|
||||||
if not _is_safe_path(req.file_path):
|
if not _is_safe_path(req.file_path):
|
||||||
@@ -85,14 +90,14 @@ def ingest_document(
|
|||||||
raise HTTPException(status_code=400, detail="不允许访问当前目录外的路径")
|
raise HTTPException(status_code=400, detail="不允许访问当前目录外的路径")
|
||||||
if not path.exists():
|
if not path.exists():
|
||||||
raise HTTPException(status_code=404, detail=f"文件不存在: {path.name}")
|
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
|
file_name = path.name
|
||||||
elif req.content:
|
elif req.content:
|
||||||
file_name = req.file_name or "untitled.md"
|
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:
|
else:
|
||||||
raise HTTPException(status_code=400, detail="需要提供 file_path 或 content")
|
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:
|
except HTTPException:
|
||||||
raise
|
raise
|
||||||
except Exception:
|
except Exception:
|
||||||
@@ -106,8 +111,9 @@ def search_documents(
|
|||||||
state: AppState = Depends(get_state),
|
state: AppState = Depends(get_state),
|
||||||
_: bool = Depends(verify_api_key),
|
_: bool = Depends(verify_api_key),
|
||||||
):
|
):
|
||||||
results = state.searcher.search(req.query, top_k=req.top_k)
|
searcher = state.get_searcher(req.collection)
|
||||||
return {"results": results}
|
results = searcher.search(req.query, top_k=req.top_k)
|
||||||
|
return {"results": results, "collection": searcher.collection_name}
|
||||||
|
|
||||||
|
|
||||||
@app.delete("/api/v1/documents/{file_name}")
|
@app.delete("/api/v1/documents/{file_name}")
|
||||||
@@ -115,8 +121,10 @@ def delete_document(
|
|||||||
file_name: str,
|
file_name: str,
|
||||||
state: AppState = Depends(get_state),
|
state: AppState = Depends(get_state),
|
||||||
_: bool = Depends(verify_api_key),
|
_: 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:
|
if not deleted:
|
||||||
raise HTTPException(status_code=404, detail=f"文档不存在: {file_name}")
|
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}
|
||||||
|
|||||||
+39
-6
@@ -12,7 +12,7 @@ logger = logging.getLogger(__name__)
|
|||||||
|
|
||||||
|
|
||||||
class AppState:
|
class AppState:
|
||||||
"""应用级共享状态."""
|
"""应用级共享状态 — db 和 embedder 全局共享, searcher/ingestor 按 collection 懒创建."""
|
||||||
|
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
config_path = os.environ.get("MD_VECTOR_CONFIG", "config.yaml")
|
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
|
"MD_VECTOR_DB_DATA_DIR", self.config.chroma.persist_dir
|
||||||
)
|
)
|
||||||
self.db = VectorDB(persist_dir=data_dir)
|
self.db = VectorDB(persist_dir=data_dir)
|
||||||
|
|
||||||
self.embedder = create_embedder(self.config.embed)
|
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
|
"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:
|
def is_healthy(self) -> dict:
|
||||||
"""真实健康检查: 验证 ChromaDB 和 Embedder 是否可用."""
|
"""真实健康检查: 验证 ChromaDB 和 Embedder 是否可用."""
|
||||||
status = {"status": "ok", "checks": {}}
|
status = {"status": "ok", "checks": {}}
|
||||||
try:
|
try:
|
||||||
count = self.searcher.collection.count()
|
count = self.db.get_or_create_collection(
|
||||||
|
self.default_collection
|
||||||
|
).count()
|
||||||
status["checks"]["chromadb"] = {"status": "ok", "count": count}
|
status["checks"]["chromadb"] = {"status": "ok", "count": count}
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
status["checks"]["chromadb"] = {"status": "error", "detail": str(e)}
|
status["checks"]["chromadb"] = {"status": "error", "detail": str(e)}
|
||||||
@@ -59,3 +87,8 @@ def get_state() -> AppState:
|
|||||||
logger.info("初始化应用状态...")
|
logger.info("初始化应用状态...")
|
||||||
_state = AppState()
|
_state = AppState()
|
||||||
return _state
|
return _state
|
||||||
|
|
||||||
|
|
||||||
|
def get_default_collection() -> str:
|
||||||
|
"""获取默认 collection 名, 供 CLI 等非 HTTP 场景使用."""
|
||||||
|
return get_state().default_collection
|
||||||
|
|||||||
Reference in New Issue
Block a user