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, }