From c2cb2562d5661d1bc999dd21d276026c7819200c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Andreas=20K=C3=B6pf?= Date: Thu, 29 Dec 2022 14:04:35 +0100 Subject: [PATCH] add db seed-data check on backend startup --- backend/main.py | 93 +++++++++++++++++++++++++++++++ backend/oasst_backend/api/deps.py | 24 ++++---- backend/oasst_backend/config.py | 1 + 3 files changed, 108 insertions(+), 10 deletions(-) diff --git a/backend/main.py b/backend/main.py index 386a495b..db1ba8d6 100644 --- a/backend/main.py +++ b/backend/main.py @@ -1,14 +1,21 @@ # -*- coding: utf-8 -*- from http import HTTPStatus from pathlib import Path +from typing import Optional import alembic.command import alembic.config import fastapi +import pydantic from loguru import logger +from oasst_backend.api.deps import get_dummy_api_client from oasst_backend.api.v1.api import api_router from oasst_backend.config import settings +from oasst_backend.database import engine from oasst_backend.exceptions import OasstError, OasstErrorCode +from oasst_backend.prompt_repository import PromptRepository +from oasst_shared.schemas import protocol as protocol_schema +from sqlmodel import Session from starlette.middleware.cors import CORSMiddleware app = fastapi.FastAPI(title=settings.PROJECT_NAME, openapi_url=f"{settings.API_V1_STR}/openapi.json") @@ -56,4 +63,90 @@ if settings.UPDATE_ALEMBIC: logger.exception("Alembic upgrade failed on startup") +if settings.USE_SEED_DATA: + + @app.on_event("startup") + def seed_data(): + class DummyPost(pydantic.BaseModel): + task_post_id: str + user_post_id: str + parent_post_id: Optional[str] + text: str + role: str + + try: + logger.info("Seed data check began") + with Session(engine) as db: + api_client = get_dummy_api_client(db) + dummy_user = protocol_schema.User(id="__dummy_user__", display_name="Dummy User", auth_method="local") + pr = PromptRepository(db=db, api_client=api_client, user=dummy_user) + + dummy_posts = [ + DummyPost( + task_post_id="de111fa8", + user_post_id="6f1d0711", + parent_post_id=None, + text="Hi!", + role="uesr", + ), + DummyPost( + task_post_id="74c381d4", + user_post_id="4a24530b", + parent_post_id="6f1d0711", + text="Hello! How can I help you?", + role="assistant", + ), + DummyPost( + task_post_id="970c437d", + user_post_id="cec432cf", + parent_post_id=None, + text="euirdteunvglfe23908230892309832098 AAAAAAAA", + role="user", + ), + DummyPost( + task_post_id="6066118e", + user_post_id="4f85f637", + parent_post_id="cec432cf", + text="Sorry, I did not understand your request and it is unclear to me what you want me to do. Could you describe it in a different way?", + role="assistant", + ), + ] + + for p in dummy_posts: + wp = pr.fetch_workpackage_by_postid(p.task_post_id) + if wp and not wp.ack: + logger.warning("Deleting unacknowledged seed data work package") + db.delete(wp) + wp = None + if not wp: + if p.parent_post_id is None: + wp = pr.store_task( + protocol_schema.InitialPromptTask(hint=""), thread_id=None, parent_post_id=None + ) + else: + print("p.parent_post_id", p.parent_post_id) + parent_post = pr.fetch_post_by_frontend_post_id(p.parent_post_id, fail_if_missing=True) + wp = pr.store_task( + protocol_schema.AssistantReplyTask( + conversation=protocol_schema.Conversation( + messages=[protocol_schema.ConversationMessage(text="dummy", is_assistant=False)] + ) + ), + thread_id=parent_post.thread_id, + parent_post_id=parent_post.id, + ) + pr.bind_frontend_post_id(wp.id, p.task_post_id) + post = pr.store_text_reply(p.text, p.task_post_id, p.user_post_id) + + logger.info( + f"Inserted: post_id: {post.id}, payload: {post.payload.payload}, parent_post_id: {post.parent_id}" + ) + else: + logger.debug(f"seed data work_package found: {wp.id}") + logger.info("Seed data check completed") + + except Exception: + logger.exception("Seed data insertion failed") + + app.include_router(api_router, prefix=settings.API_V1_STR) diff --git a/backend/oasst_backend/api/deps.py b/backend/oasst_backend/api/deps.py index bdbd83eb..244e55c7 100644 --- a/backend/oasst_backend/api/deps.py +++ b/backend/oasst_backend/api/deps.py @@ -33,6 +33,19 @@ async def get_api_key( return api_key_header +def get_dummy_api_client(db: Session) -> ApiClient: + # make sure that a dummy api key exits in db (foreign key references) + ANY_API_KEY_ID = UUID("00000000-1111-2222-3333-444444444444") + api_client: ApiClient = db.query(ApiClient).filter(ApiClient.id == ANY_API_KEY_ID).first() + if api_client is None: + token = token_hex(32) + logger.info(f"ANY_API_KEY missing, inserting api_key: {token}") + api_client = ApiClient(id=ANY_API_KEY_ID, api_key=token, description="ANY_API_KEY, random token") + db.add(api_client) + db.commit() + return api_client + + def api_auth( api_key: APIKey, db: Session, @@ -40,16 +53,7 @@ def api_auth( if api_key or settings.DEBUG_SKIP_API_KEY_CHECK: if settings.DEBUG_SKIP_API_KEY_CHECK or settings.DEBUG_ALLOW_ANY_API_KEY: - # make sure that a dummy api key exits in db (foreign key references) - ANY_API_KEY_ID = UUID("00000000-1111-2222-3333-444444444444") - api_client: ApiClient = db.query(ApiClient).filter(ApiClient.id == ANY_API_KEY_ID).first() - if api_client is None: - token = token_hex(32) - logger.info(f"ANY_API_KEY missing, inserting api_key: {token}") - api_client = ApiClient(id=ANY_API_KEY_ID, api_key=token, description="ANY_API_KEY, random token") - db.add(api_client) - db.commit() - return api_client + return get_dummy_api_client() api_client = db.query(ApiClient).filter(ApiClient.api_key == api_key).first() if api_client is not None and api_client.enabled: diff --git a/backend/oasst_backend/config.py b/backend/oasst_backend/config.py index 96d6021e..ee18c180 100644 --- a/backend/oasst_backend/config.py +++ b/backend/oasst_backend/config.py @@ -33,6 +33,7 @@ class Settings(BaseSettings): BACKEND_CORS_ORIGINS: List[AnyHttpUrl] = [] UPDATE_ALEMBIC: bool = True + USE_SEED_DATA: bool = True @validator("BACKEND_CORS_ORIGINS", pre=True) def assemble_cors_origins(cls, v: Union[str, List[str]]) -> Union[List[str], str]: