diff --git a/backend/oasst_backend/api/v1/api.py b/backend/oasst_backend/api/v1/api.py index b39d8d59..2286a1ac 100644 --- a/backend/oasst_backend/api/v1/api.py +++ b/backend/oasst_backend/api/v1/api.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- from fastapi import APIRouter -from oasst_backend.api.v1 import frontend_messages, frontend_users, messages, tasks, text_labels, users +from oasst_backend.api.v1 import frontend_messages, frontend_users, messages, stats, tasks, text_labels, users api_router = APIRouter() api_router.include_router(tasks.router, prefix="/tasks", tags=["tasks"]) @@ -9,3 +9,4 @@ api_router.include_router(messages.router, prefix="/messages", tags=["messages"] api_router.include_router(frontend_messages.router, prefix="/frontend_messages", tags=["frontend_messages"]) api_router.include_router(users.router, prefix="/users", tags=["users"]) api_router.include_router(frontend_users.router, prefix="/frontend_users", tags=["frontend_users"]) +api_router.include_router(stats.router, prefix="/stats", tags=["stats"]) diff --git a/backend/oasst_backend/api/v1/stats.py b/backend/oasst_backend/api/v1/stats.py new file mode 100644 index 00000000..789e53f9 --- /dev/null +++ b/backend/oasst_backend/api/v1/stats.py @@ -0,0 +1,18 @@ +# -*- coding: utf-8 -*- +from fastapi import APIRouter, Depends +from oasst_backend.api import deps +from oasst_backend.models import ApiClient +from oasst_backend.prompt_repository import PromptRepository +from oasst_shared.schemas import protocol +from sqlmodel import Session + +router = APIRouter() + + +@router.get("/") +def get_message_stats( + db: Session = Depends(deps.get_db), + api_client: ApiClient = Depends(deps.get_trusted_api_client), +): + pr = PromptRepository(db, api_client, None) + return protocol.SystemStats(**pr.get_stats()) diff --git a/backend/oasst_backend/prompt_repository.py b/backend/oasst_backend/prompt_repository.py index 05902b5a..8df17347 100644 --- a/backend/oasst_backend/prompt_repository.py +++ b/backend/oasst_backend/prompt_repository.py @@ -663,3 +663,15 @@ class PromptRepository: parent_ids = self.db.execute(query).scalars().all() self.db.commit() + + def get_stats(self): + deleted = self.db.query(Post.deleted, func.count()).group_by(Post.deleted) + nthreads = self.db.query(None, func.count(Post.id)).filter(Post.parent_id.is_(None)) + query = deleted.union_all(nthreads) + result = {k: v for k, v in query.all()} + return { + "all": result.get(True, 0) + result.get(False, 0), + "active": result.get(False, 0), + "deleted": result.get(True, 0), + "threads": result.get(None, 0), + } diff --git a/oasst-shared/oasst_shared/schemas/protocol.py b/oasst-shared/oasst_shared/schemas/protocol.py index 028154c9..a2a12ddc 100644 --- a/oasst-shared/oasst_shared/schemas/protocol.py +++ b/oasst-shared/oasst_shared/schemas/protocol.py @@ -271,3 +271,10 @@ class TextLabels(BaseModel): if not (0 <= value <= 1): raise ValueError(f"Label values must be between 0 and 1, got {value} for {key}.") return v + + +class SystemStats(BaseModel): + all: int = 0 + active: int = 0 + deleted: int = 0 + threads: int = 0