Spaces:
Running
Running
File size: 10,490 Bytes
0752cd0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 | from __future__ import annotations
import asyncio
import os
import shutil
from pathlib import Path
from typing import Any, Dict, List, Optional
from fastapi import APIRouter, Depends, HTTPException, Query, status
from pydantic import BaseModel, Field
from app.auth import get_current_user
from app.compare.collection_registry import clear_collection_registry
from app.models.user import User
from app.services.pipeline_manager import PipelineManager
try:
from chromadb import PersistentClient
except Exception: # pragma: no cover - chromadb should be installed in runtime
PersistentClient = None
router = APIRouter(prefix="/api/admin", tags=["admin"])
def _is_admin_user(user: User) -> bool:
if getattr(user, "is_admin", False):
return True
seed_username = os.getenv("AUTH_SEED_USERNAME", "admin")
seed_email = os.getenv("AUTH_SEED_EMAIL", "admin@local")
return user.username == seed_username or user.email == seed_email
def require_admin(current_user: User = Depends(get_current_user)) -> User:
if not _is_admin_user(current_user):
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Admin access required")
return current_user
def _project_root() -> Path:
return Path(__file__).resolve().parents[2]
def _storage_roots() -> List[Path]:
roots = [
Path(os.getenv("CHROMA_PERSIST_DIR", str(_project_root() / "chroma_db"))),
_project_root() / "chroma_db",
_project_root() / "chroma_store",
_project_root() / "chroma_data",
]
deduped: List[Path] = []
seen = set()
for root in roots:
resolved = root.resolve()
if resolved not in seen:
seen.add(resolved)
deduped.append(resolved)
return deduped
def _collection_client(root: Path):
if PersistentClient is None:
raise HTTPException(status_code=500, detail="Chroma client is not available")
root.mkdir(parents=True, exist_ok=True)
return PersistentClient(path=str(root))
def _clear_collections_in_root(root: Path) -> List[str]:
"""Delete all collections from a specific root using Chroma APIs first."""
try:
client = _collection_client(root)
collections = client.list_collections()
except Exception:
return []
deleted: List[str] = []
for collection in collections:
name = getattr(collection, "name", collection if isinstance(collection, str) else None)
if not name:
continue
try:
client.delete_collection(name=name)
deleted.append(str(name))
except Exception:
continue
return deleted
def _clear_runtime_caches() -> None:
"""Release in-memory references that can keep Chroma files/collections active."""
try:
clear_collection_registry()
except Exception:
pass
try:
PipelineManager.clear_cache()
except Exception:
pass
def _list_root_collections(root: Path) -> List[Dict[str, Any]]:
try:
client = _collection_client(root)
collections = client.list_collections()
except Exception:
return []
summaries: List[Dict[str, Any]] = []
for collection in collections:
collection_obj = collection
if isinstance(collection, str):
try:
collection_obj = client.get_collection(name=collection)
except Exception:
collection_obj = None
if collection_obj is None:
continue
try:
count = int(collection_obj.count())
except Exception:
count = 0
sample_docs: List[Dict[str, Any]] = []
try:
sample = collection_obj.get(limit=3, include=["documents", "metadatas"])
ids = sample.get("ids", []) if isinstance(sample, dict) else []
documents = sample.get("documents", []) if isinstance(sample, dict) else []
metadatas = sample.get("metadatas", []) if isinstance(sample, dict) else []
for idx, sample_id in enumerate(ids[:3]):
sample_docs.append({
"id": sample_id,
"document": documents[idx] if idx < len(documents) else None,
"metadata": metadatas[idx] if idx < len(metadatas) else {},
})
except Exception:
sample_docs = []
summaries.append({
"name": getattr(collection_obj, "name", collection if isinstance(collection, str) else "unknown"),
"count": count,
"metadata": getattr(collection_obj, "metadata", {}) or {},
"samples": sample_docs,
})
return summaries
class ChromaCollectionDetail(BaseModel):
name: str
count: int = 0
metadata: Dict[str, Any] = Field(default_factory=dict)
samples: List[Dict[str, Any]] = Field(default_factory=list)
class ChromaRootDetail(BaseModel):
root_path: str
collections: List[ChromaCollectionDetail] = Field(default_factory=list)
class ChromaDeleteResponse(BaseModel):
status: str
deleted: List[str] = Field(default_factory=list)
@router.get("/chroma", response_model=List[ChromaRootDetail])
async def list_chroma_roots(current_user: User = Depends(require_admin)):
_ = current_user
roots = []
for root in _storage_roots():
collections = await asyncio.to_thread(_list_root_collections, root)
roots.append(
ChromaRootDetail(
root_path=str(root),
collections=[ChromaCollectionDetail(**collection) for collection in collections],
)
)
return roots
@router.get("/chroma/collections/{collection_name}", response_model=List[ChromaRootDetail])
async def view_collection(collection_name: str, current_user: User = Depends(require_admin)):
_ = current_user
roots: List[ChromaRootDetail] = []
for root in _storage_roots():
collections = [collection for collection in await asyncio.to_thread(_list_root_collections, root) if collection["name"] == collection_name]
if collections:
roots.append(ChromaRootDetail(root_path=str(root), collections=[ChromaCollectionDetail(**collection) for collection in collections]))
return roots
@router.delete("/chroma/collections/{collection_name}", response_model=ChromaDeleteResponse)
async def delete_collection(
collection_name: str,
root_path: Optional[str] = Query(None),
current_user: User = Depends(require_admin),
):
_ = current_user
deleted: List[str] = []
roots = _storage_roots() if root_path is None else [Path(root_path).resolve()]
for root in roots:
try:
await asyncio.to_thread(lambda: _collection_client(root).delete_collection(name=collection_name))
deleted.append(f"{root}:{collection_name}")
except Exception:
continue
if not deleted:
raise HTTPException(status_code=404, detail="Collection not found")
return ChromaDeleteResponse(status="success", deleted=deleted)
@router.delete("/chroma/root", response_model=ChromaDeleteResponse)
async def clear_root(
root_path: Optional[str] = Query(None),
current_user: User = Depends(require_admin),
):
_ = current_user
roots = [Path(root_path).resolve()] if root_path else _storage_roots()
deleted: List[str] = []
await asyncio.to_thread(_clear_runtime_caches)
for target_root in roots:
try:
await asyncio.to_thread(_clear_collections_in_root, target_root)
except Exception:
pass
if target_root.exists():
try:
await asyncio.to_thread(shutil.rmtree, target_root)
except Exception:
# If filesystem removal is blocked by locks, keep folder but collections
# are already deleted through Chroma API.
pass
return ChromaDeleteResponse(status="success", deleted=deleted)
@router.get("/db-status")
async def db_status(current_user: User = Depends(require_admin)):
"""Diagnostic endpoint: shows storage paths, document counts, and current user's documents.
Helps diagnose 'documents invisible after upload' issues caused by storage
path mismatches or DB isolation problems.
"""
from sqlalchemy import func, select, text
from sqlalchemy.ext.asyncio import AsyncSession
from app.database import AsyncSessionLocal
from app.models.document import Document
from app.models.user import User as UserModel
async with AsyncSessionLocal() as db:
# Total documents in DB
total_docs = (await db.execute(select(func.count()).select_from(Document))).scalar()
# Documents per user
per_user_rows = (await db.execute(
select(Document.user_id, func.count(Document.id).label("count"))
.group_by(Document.user_id)
)).all()
per_user = [{"user_id": str(r[0]), "count": r[1]} for r in per_user_rows]
# Current user's documents (most recent 10)
my_docs_rows = (await db.execute(
select(Document.id, Document.filename, Document.file_type, Document.upload_date)
.where(Document.user_id == current_user.id)
.order_by(Document.upload_date.desc())
.limit(10)
)).all()
my_docs = [
{"id": str(r[0]), "filename": r[1], "file_type": r[2], "upload_date": str(r[3])}
for r in my_docs_rows
]
# All users
all_users = (await db.execute(select(UserModel.id, UserModel.username))).all()
users = [{"id": str(r[0]), "username": r[1]} for r in all_users]
return {
"storage": {
"chroma_persist_dir": os.getenv("CHROMA_PERSIST_DIR", "(not set)"),
"upload_dir": os.getenv("UPLOAD_DIR", "(not set)"),
"database_url": os.getenv("DATABASE_URL", "(not set)"),
},
"current_user": {
"id": str(current_user.id),
"username": current_user.username,
},
"all_users": users,
"total_documents": total_docs,
"documents_per_user": per_user,
"my_recent_documents": my_docs,
}
|