add periodically updated cached message stats (LAION-AI#1464)

* add periodically updated cached message stats

Author: andreaskoepf

backend/alembic/versions/2023_02_11_1030-ba40d055714a_add_cached_stats.py

@@ -0,0 +1,37 @@
+"""add cached_stats
+
+Revision ID: ba40d055714a
+Revises: caee1e8ee0bc
+Create Date: 2023-02-11 10:30:21.996198
+
+"""
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "ba40d055714a"
+down_revision = "caee1e8ee0bc"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "cached_stats",
+        sa.Column("name", sqlmodel.sql.sqltypes.AutoString(length=128), nullable=False),
+        sa.Column(
+            "modified_date", sa.DateTime(timezone=True), server_default=sa.text("CURRENT_TIMESTAMP"), nullable=False
+        ),
+        sa.Column("stats", postgresql.JSONB(astext_type=sa.Text()), nullable=False),
+        sa.PrimaryKeyConstraint("name"),
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table("cached_stats")
+    # ### end Alembic commands ###

backend/main.py

@@ -15,6 +15,7 @@
 from oasst_backend.api.deps import api_auth, create_api_client
 from oasst_backend.api.v1.api import api_router
 from oasst_backend.api.v1.utils import prepare_conversation
+from oasst_backend.cached_stats_repository import CachedStatsRepository
 from oasst_backend.config import settings
 from oasst_backend.database import engine
 from oasst_backend.models import message_tree_state
@@ -334,6 +335,17 @@ def cronjob_delete_expired_tasks(session: Session) -> None:
     delete_expired_tasks(session)
 
 
+@app.on_event("startup")
+@repeat_every(seconds=60 * settings.CACHED_STATS_UPDATE_INTERVAL, wait_first=True)
+@managed_tx_function(auto_commit=CommitMode.COMMIT)
+def update_cached_stats(session: Session) -> None:
+    try:
+        csr = CachedStatsRepository(session)
+        csr.update_all_cached_stats()
+    except Exception:
+        logger.exception("Error during cached stats update")
+
+
 app.include_router(api_router, prefix=settings.API_V1_STR)
 
 

backend/oasst_backend/api/v1/stats.py

@@ -1,10 +1,13 @@
 from fastapi import APIRouter, Depends
 from oasst_backend.api import deps
+from oasst_backend.cached_stats_repository import CachedStatsRepository
 from oasst_backend.models import ApiClient
 from oasst_backend.prompt_repository import PromptRepository
 from oasst_backend.tree_manager import TreeManager, TreeManagerStats, TreeMessageCountStats
+from oasst_backend.utils.database_utils import CommitMode, managed_tx_function
 from oasst_shared.schemas import protocol
 from sqlmodel import Session
+from starlette.status import HTTP_204_NO_CONTENT
 
 router = APIRouter()
 
@@ -47,3 +50,38 @@ def get_tree_manager__stats(
     pr = PromptRepository(db, api_client)
     tm = TreeManager(db, pr)
     return tm.stats()
+
+
+@router.get("/cached/{name}", response_model=protocol.CachedStatsResponse)
+def get_cached_stats(
+    *,
+    name: protocol.CachedStatsName,
+    db: Session = Depends(deps.get_db),
+    api_client: ApiClient = Depends(deps.get_api_client),
+):
+    csr = CachedStatsRepository(db)
+    return csr.get_stats(name)
+
+
+@router.get("/cached", response_model=protocol.AllCachedStatsResponse)
+def get_cached_stats_all(
+    *,
+    db: Session = Depends(deps.get_db),
+    api_client: ApiClient = Depends(deps.get_api_client),
+):
+    csr = CachedStatsRepository(db)
+    return csr.get_stats_all()
+
+
+@router.post("/cached/update", response_model=None, status_code=HTTP_204_NO_CONTENT)
+def update_cached_stats(
+    *,
+    db: Session = Depends(deps.get_db),
+    api_client: ApiClient = Depends(deps.get_trusted_api_client),
+):
+    @managed_tx_function(CommitMode.COMMIT)
+    def update_tx(db: deps.Session) -> None:
+        csr = CachedStatsRepository(db)
+        csr.update_all_cached_stats()
+
+    update_tx()

backend/oasst_backend/cached_stats_repository.py

@@ -0,0 +1,103 @@
+from oasst_backend.models import CachedStats, Message, MessageTreeState, User
+from oasst_shared.exceptions.oasst_api_error import OasstError, OasstErrorCode
+from oasst_shared.schemas.protocol import AllCachedStatsResponse, CachedStatsName, CachedStatsResponse
+from oasst_shared.utils import log_timing, utcnow
+from sqlalchemy.orm.attributes import flag_modified
+from sqlmodel import Session, func, not_
+
+
+def row_to_dict(r) -> dict:
+    return {k: r[k] for k in r.keys()}
+
+
+class CachedStatsRepository:
+    def __init__(self, db: Session):
+        self.db = db
+
+    def qry_human_messages_by_lang(self) -> dict[str, int]:
+        qry = (
+            self.db.query(Message.lang, func.count(Message.id).label("count"))
+            .filter(not_(Message.deleted), Message.review_result, not_(Message.synthetic))
+            .group_by(Message.lang)
+        )
+        return {r["lang"]: r["count"] for r in qry}
+
+    def qry_human_messages_by_role(self) -> dict[str, int]:
+        qry = (
+            self.db.query(Message.role, func.count(Message.id).label("count"))
+            .filter(not_(Message.deleted), Message.review_result, not_(Message.synthetic))
+            .group_by(Message.role)
+        )
+        return {r["role"]: r["count"] for r in qry}
+
+    def qry_message_trees_by_state(self) -> dict[str, int]:
+        qry = self.db.query(
+            MessageTreeState.state, func.count(MessageTreeState.message_tree_id).label("count")
+        ).group_by(MessageTreeState.state)
+        return {r["state"]: r["count"] for r in qry}
+
+    def qry_message_trees_states_by_lang(self) -> list:
+        qry = (
+            self.db.query(
+                Message.lang, MessageTreeState.state, func.count(MessageTreeState.message_tree_id).label("count")
+            )
+            .select_from(MessageTreeState)
+            .join(Message, MessageTreeState.message_tree_id == Message.id)
+            .group_by(MessageTreeState.state, Message.lang)
+            .order_by(Message.lang, MessageTreeState.state)
+        )
+        return [row_to_dict(r) for r in qry]
+
+    def qry_users_accepted_tos(self) -> dict[str, int]:
+        qry = self.db.query(func.count(User.id)).filter(User.enabled, User.tos_acceptance_date.is_not(None))
+        return {"count": qry.scalar()}
+
+    @log_timing(level="INFO")
+    def update_all_cached_stats(self):
+        v = self.qry_human_messages_by_lang()
+        self._insert_cached_stats(CachedStatsName.human_messages_by_lang, v)
+
+        v = self.qry_human_messages_by_role()
+        self._insert_cached_stats(CachedStatsName.human_messages_by_role, v)
+
+        v = self.qry_message_trees_by_state()
+        self._insert_cached_stats(CachedStatsName.message_trees_by_state, v)
+
+        v = self.qry_message_trees_states_by_lang()
+        self._insert_cached_stats(CachedStatsName.message_trees_states_by_lang, v)
+
+        v = self.qry_users_accepted_tos()
+        self._insert_cached_stats(CachedStatsName.users_accepted_tos, v)
+
+    def _insert_cached_stats(self, name: CachedStatsName, stats: dict | list):
+        row: CachedStats | None = self.db.query(CachedStats).filter(CachedStats.name == name).one_or_none()
+        if row:
+            row.modified_date = utcnow()
+            row.stats = stats
+            flag_modified(row, "stats")
+        else:
+            row = CachedStats(name=name, modified_date=utcnow(), stats=stats)
+        self.db.add(row)
+
+    def get_stats(self, name: CachedStatsName) -> CachedStatsResponse:
+        row: CachedStats | None = self.db.query(CachedStats).filter(CachedStats.name == name).one_or_none()
+        if not row:
+            raise OasstError(f"Cached stats '{name.value}' not found.", OasstErrorCode.CACHED_STATS_NOT_AVAILABLE)
+        return CachedStatsResponse(name=row.name, last_updated=row.modified_date, stats=row.stats)
+
+    def get_stats_all(self) -> AllCachedStatsResponse:
+        by_name: dict[CachedStatsName, CachedStatsResponse] = {}
+        qry = self.db.query(CachedStats)
+        for row in qry:
+            by_name[row.name] = CachedStatsResponse(name=row.name, last_updated=row.modified_date, stats=row.stats)
+        return AllCachedStatsResponse(stats_by_name=by_name)
+
+
+if __name__ == "__main__":
+    # from oasst_backend.api.deps import create_api_client
+    from oasst_backend.database import engine
+
+    with Session(engine) as db:
+        csr = CachedStatsRepository(db)
+        csr.update_all_cached_stats()()
+        db.commit()

backend/oasst_backend/config.py

@@ -243,6 +243,8 @@ def validate_user_stats_intervals(cls, v: int):
             raise ValueError(v)
         return v
 
+    CACHED_STATS_UPDATE_INTERVAL: int = 60  # minutes
+
     RATE_LIMIT_TASK_USER_TIMES: int = 60
     RATE_LIMIT_TASK_USER_MINUTES: int = 5
     RATE_LIMIT_TASK_API_TIMES: int = 10_000

backend/oasst_backend/models/__init__.py

@@ -1,4 +1,5 @@
 from .api_client import ApiClient
+from .cached_stats import CachedStats
 from .flagged_message import FlaggedMessage
 from .journal import Journal, JournalIntegration
 from .message import Message
@@ -30,4 +31,5 @@
     "MessageEmoji",
     "TrollStats",
     "FlaggedMessage",
+    "CachedStats",
 ]

backend/oasst_backend/models/cached_stats.py

@@ -0,0 +1,17 @@
+from datetime import datetime
+
+import sqlalchemy as sa
+import sqlalchemy.dialects.postgresql as pg
+from sqlmodel import AutoString, Field, SQLModel
+
+
+class CachedStats(SQLModel, table=True):
+    __tablename__ = "cached_stats"
+
+    name: str = Field(sa_column=sa.Column(AutoString(length=128), primary_key=True))
+
+    modified_date: datetime | None = Field(
+        sa_column=sa.Column(sa.DateTime(timezone=True), nullable=False, server_default=sa.func.current_timestamp())
+    )
+
+    stats: dict | list | None = Field(None, sa_column=sa.Column(pg.JSONB, nullable=False))

oasst-shared/oasst_shared/exceptions/oasst_api_error.py

@@ -86,6 +86,8 @@ class OasstErrorCode(IntEnum):
 
     EMOJI_OP_UNSUPPORTED = 5000
 
+    CACHED_STATS_NOT_AVAILABLE = 6000
+
 
 class OasstError(Exception):
     """Base class for Open-Assistant exceptions."""

oasst-shared/oasst_shared/schemas/protocol.py

@@ -571,3 +571,21 @@ class CreateFrontendUserRequest(User):
     enabled: bool = True
     tos_acceptance: Optional[bool] = None
     notes: Optional[str] = None
+
+
+class CachedStatsName(str, enum.Enum):
+    human_messages_by_lang = "human_messages_by_lang"
+    human_messages_by_role = "human_messages_by_role"
+    message_trees_by_state = "message_trees_by_state"
+    message_trees_states_by_lang = "message_trees_states_by_lang"
+    users_accepted_tos = "users_accepted_tos"
+
+
+class CachedStatsResponse(BaseModel):
+    name: CachedStatsName | str
+    last_updated: datetime
+    stats: dict | list
+
+
+class AllCachedStatsResponse(BaseModel):
+    stats_by_name: dict[CachedStatsName | str, CachedStatsResponse]