add db seed-data check on backend startup

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)

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:

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]: