import json from datetime import datetime from http import HTTPStatus from math import ceil from pathlib import Path from typing import Optional import alembic.command import alembic.config import fastapi import redis.asyncio as redis from fastapi_limiter import FastAPILimiter from fastapi_utils.tasks import repeat_every from loguru import logger 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.config import settings from oasst_backend.database import engine from oasst_backend.models import message_tree_state from oasst_backend.prompt_repository import PromptRepository, UserRepository from oasst_backend.task_repository import TaskRepository, delete_expired_tasks from oasst_backend.tree_manager import TreeManager from oasst_backend.user_repository import User from oasst_backend.user_stats_repository import UserStatsRepository, UserStatsTimeFrame from oasst_backend.utils.database_utils import CommitMode, managed_tx_function from oasst_shared.exceptions import OasstError, OasstErrorCode from oasst_shared.schemas import protocol as protocol_schema from oasst_shared.utils import utcnow from pydantic import BaseModel from sqlmodel import Session, select from starlette.middleware.cors import CORSMiddleware app = fastapi.FastAPI(title=settings.PROJECT_NAME, openapi_url=f"{settings.API_V1_STR}/openapi.json") startup_time: datetime = utcnow() @app.exception_handler(OasstError) async def oasst_exception_handler(request: fastapi.Request, ex: OasstError): logger.error(f"{request.method} {request.url} failed: {repr(ex)}") return fastapi.responses.JSONResponse( status_code=int(ex.http_status_code), content=protocol_schema.OasstErrorResponse( message=ex.message, error_code=OasstErrorCode(ex.error_code), ).dict(), ) @app.exception_handler(Exception) async def unhandled_exception_handler(request: fastapi.Request, ex: Exception): logger.exception(f"{request.method} {request.url} failed [UNHANDLED]: {repr(ex)}") status = HTTPStatus.INTERNAL_SERVER_ERROR return fastapi.responses.JSONResponse( status_code=status.value, content={"message": status.name, "error_code": OasstErrorCode.GENERIC_ERROR} ) # Set all CORS enabled origins if settings.BACKEND_CORS_ORIGINS: app.add_middleware( CORSMiddleware, allow_origins=[str(origin) for origin in settings.BACKEND_CORS_ORIGINS], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) if settings.UPDATE_ALEMBIC: @app.on_event("startup") def alembic_upgrade(): logger.info("Attempting to upgrade alembic on startup") try: alembic_ini_path = Path(__file__).parent / "alembic.ini" alembic_cfg = alembic.config.Config(str(alembic_ini_path)) alembic_cfg.set_main_option("sqlalchemy.url", settings.DATABASE_URI) alembic.command.upgrade(alembic_cfg, "head") logger.info("Successfully upgraded alembic on startup") except Exception: logger.exception("Alembic upgrade failed on startup") if settings.OFFICIAL_WEB_API_KEY: @app.on_event("startup") def create_official_web_api_client(): with Session(engine) as session: try: api_auth(settings.OFFICIAL_WEB_API_KEY, db=session) except OasstError: logger.info("Creating official web API client") create_api_client( session=session, api_key=settings.OFFICIAL_WEB_API_KEY, description="The official web client for the OASST backend.", frontend_type="web", trusted=True, ) if settings.RATE_LIMIT: @app.on_event("startup") async def connect_redis(): async def http_callback(request: fastapi.Request, response: fastapi.Response, pexpire: int): """Error callback function when too many requests""" expire = ceil(pexpire / 1000) raise OasstError( f"Too Many Requests. Retry After {expire} seconds.", OasstErrorCode.TOO_MANY_REQUESTS, HTTPStatus.TOO_MANY_REQUESTS, ) try: redis_client = redis.from_url( f"redis://{settings.REDIS_HOST}:{settings.REDIS_PORT}/0", encoding="utf-8", decode_responses=True ) logger.info(f"Connected to {redis_client=}") await FastAPILimiter.init(redis_client, http_callback=http_callback) except Exception: logger.exception("Failed to establish Redis connection") if settings.DEBUG_USE_SEED_DATA: @app.on_event("startup") @managed_tx_function(auto_commit=CommitMode.COMMIT) def create_seed_data(session: Session): class DummyMessage(BaseModel): task_message_id: str user_message_id: str parent_message_id: Optional[str] text: str lang: Optional[str] role: str tree_state: Optional[message_tree_state.State] if not settings.OFFICIAL_WEB_API_KEY: raise ValueError("Cannot use seed data without OFFICIAL_WEB_API_KEY") try: logger.info("Seed data check began") api_client = api_auth(settings.OFFICIAL_WEB_API_KEY, db=session) dummy_user = protocol_schema.User(id="__dummy_user__", display_name="Dummy User", auth_method="local") ur = UserRepository(db=session, api_client=api_client) tr = TaskRepository(db=session, api_client=api_client, client_user=dummy_user, user_repository=ur) ur.update_user(tr.user_id, enabled=True, show_on_leaderboard=False, tos_acceptance=True) pr = PromptRepository( db=session, api_client=api_client, client_user=dummy_user, user_repository=ur, task_repository=tr ) tm = TreeManager(session, pr) with open(settings.DEBUG_USE_SEED_DATA_PATH) as f: dummy_messages_raw = json.load(f) dummy_messages = [DummyMessage(**dm) for dm in dummy_messages_raw] for msg in dummy_messages: task = tr.fetch_task_by_frontend_message_id(msg.task_message_id) if task and not task.ack: logger.warning("Deleting unacknowledged seed data task") session.delete(task) task = None if not task: if msg.parent_message_id is None: task = tr.store_task( protocol_schema.InitialPromptTask(hint=""), message_tree_id=None, parent_message_id=None ) else: parent_message = pr.fetch_message_by_frontend_message_id( msg.parent_message_id, fail_if_missing=True ) conversation_messages = pr.fetch_message_conversation(parent_message) conversation = prepare_conversation(conversation_messages) if msg.role == "assistant": task = tr.store_task( protocol_schema.AssistantReplyTask(conversation=conversation), message_tree_id=parent_message.message_tree_id, parent_message_id=parent_message.id, ) else: task = tr.store_task( protocol_schema.PrompterReplyTask(conversation=conversation), message_tree_id=parent_message.message_tree_id, parent_message_id=parent_message.id, ) tr.bind_frontend_message_id(task.id, msg.task_message_id) message = pr.store_text_reply( msg.text, msg.lang, msg.task_message_id, msg.user_message_id, review_count=5, review_result=True, check_tree_state=False, check_duplicate=False, ) if message.parent_id is None: tm._insert_default_state( root_message_id=message.id, state=msg.tree_state or message_tree_state.State.GROWING ) session.flush() logger.info( f"Inserted: message_id: {message.id}, payload: {message.payload.payload}, parent_message_id: {message.parent_id}" ) else: logger.debug(f"seed data task found: {task.id}") logger.info("Seed data check completed") except Exception: logger.exception("Seed data insertion failed") @app.on_event("startup") def ensure_tree_states(): try: logger.info("Startup: TreeManager.ensure_tree_states()") with Session(engine) as db: api_client = api_auth(settings.OFFICIAL_WEB_API_KEY, db=db) tm = TreeManager(db, PromptRepository(db, api_client=api_client)) tm.ensure_tree_states() except Exception: logger.exception("TreeManager.ensure_tree_states() failed.") @app.on_event("startup") @repeat_every(seconds=60 * settings.USER_STATS_INTERVAL_DAY, wait_first=False) @managed_tx_function(auto_commit=CommitMode.COMMIT) def update_leader_board_day(session: Session) -> None: try: usr = UserStatsRepository(session) usr.update_stats(time_frame=UserStatsTimeFrame.day) except Exception: logger.exception("Error during leaderboard update (daily)") @app.on_event("startup") @repeat_every(seconds=60 * settings.USER_STATS_INTERVAL_WEEK, wait_first=False) @managed_tx_function(auto_commit=CommitMode.COMMIT) def update_leader_board_week(session: Session) -> None: try: usr = UserStatsRepository(session) usr.update_stats(time_frame=UserStatsTimeFrame.week) except Exception: logger.exception("Error during user states update (weekly)") @app.on_event("startup") @repeat_every(seconds=60 * settings.USER_STATS_INTERVAL_MONTH, wait_first=False) @managed_tx_function(auto_commit=CommitMode.COMMIT) def update_leader_board_month(session: Session) -> None: try: usr = UserStatsRepository(session) usr.update_stats(time_frame=UserStatsTimeFrame.month) except Exception: logger.exception("Error during user states update (monthly)") @app.on_event("startup") @repeat_every(seconds=60 * settings.USER_STATS_INTERVAL_TOTAL, wait_first=False) @managed_tx_function(auto_commit=CommitMode.COMMIT) def update_leader_board_total(session: Session) -> None: try: usr = UserStatsRepository(session) usr.update_stats(time_frame=UserStatsTimeFrame.total) except Exception: logger.exception("Error during user states update (total)") @app.on_event("startup") @repeat_every(seconds=60 * 60 * settings.USER_STREAK_UPDATE_INTERVAL, wait_first=False) def update_user_streak(session: Session) -> None: try: current_time = utcnow() timedelta = current_time - startup_time if timedelta.days > 0: # Update only greater than 24 hours . Do nothing logger.debug("Process timedelta greater than 24h") statement = select(User) result = session.exec(statement).all() if result is not None: for user in result: last_activity_date = user.last_activity_date streak_last_day_date = user.streak_last_day_date # set NULL streak_days to 0 if user.streak_days is None: user.streak_days = 0 # if the user had completed a task if last_activity_date is not None: lastactitvitydelta = current_time - last_activity_date # if the user missed consecutive days of completing a task # reset the streak_days to 0 and set streak_last_day_date to the current_time if lastactitvitydelta.days > 1 or user.streak_days is None: user.streak_days = 0 user.streak_last_day_date = current_time # streak_last_day_date has a current timestamp in DB. Idealy should not be NULL. if streak_last_day_date is not None: streak_delta = current_time - streak_last_day_date # if user completed tasks on consecutive days then increment the streak days # update the streak_last_day_date to current time for the next calculation if streak_delta.days > 0: user.streak_days += 1 user.streak_last_day_date = current_time session.add(user) session.commit() else: logger.debug("Not yet 24hours since the process started! ...") except Exception as e: logger.error(str(e)) return @app.on_event("startup") @repeat_every(seconds=60 * 60) # 1 hour @managed_tx_function(auto_commit=CommitMode.COMMIT) def cronjob_delete_expired_tasks(session: Session) -> None: delete_expired_tasks(session) app.include_router(api_router, prefix=settings.API_V1_STR) def get_openapi_schema(): return json.dumps(app.openapi()) def retry_scoring_failed_message_trees(): try: logger.info("TreeManager.retry_scoring_failed_message_trees()") with Session(engine) as db: api_client = api_auth(settings.OFFICIAL_WEB_API_KEY, db=db) pr = PromptRepository(db=db, api_client=api_client) tm = TreeManager(db, pr) tm.retry_scoring_failed_message_trees() except Exception: logger.exception("TreeManager.retry_scoring_failed_message_trees() failed.") def main(): # Importing here so we don't import packages unnecessarily if we're # importing main as a module. import argparse import uvicorn parser = argparse.ArgumentParser() parser.add_argument( "--print-openapi-schema", default=False, help="Dumps the openapi schema to stdout", action="store_true", ) parser.add_argument("--host", help="The host to run the server", default="0.0.0.0") parser.add_argument("--port", help="The port to run the server", default=8080) parser.add_argument( "--retry-scoring", default=False, help="Retry scoring failed message trees", action="store_true", ) args = parser.parse_args() if args.print_openapi_schema: print(get_openapi_schema()) if args.retry_scoring: retry_scoring_failed_message_trees() if not (args.print_openapi_schema or args.retry_scoring): uvicorn.run(app, host=args.host, port=args.port) if __name__ == "__main__": main()