Merge pull request #185 from PrefectHQ/database-upgrade

docs/guide/concepts/tui.md

@@ -1,4 +1,4 @@
-# 🖥️ ChatTUI
+# 🖥️ TUI
 ![](../../img/tui/colorful_fruit.png)
 
 !!! tip "Features"
@@ -97,6 +97,14 @@ You can set your OpenAI API key from the TUI by pressing the Settings button. It
 
 ![](../../img/tui/settings.png)
 
+## Upgrading the database
+
+When new versions of Marvin are released, they may require your database to be upgraded in order to handle the new features or enhancements they contain. Therefore, whenever Marvin starts, it checks to see if your database is up-to-date and prints a helpful warning if it isn't. The TUI can go a step farther and upgrade your database automatically. If it is possible to do so, you will see a screen like the one below.
+
+![](../../img/tui/database_upgrade.png)
+
+!!! tip "Manually upgrading the database"
+    We have tried to prevent the TUI from accessing the database until after it runs the upgrade check and shows the warning screen. However, if for some reason it is unable to run the check properly, it may crash. In this case you can manually upgrade the database by running `marvin database upgrade` in your terminal.
 ## Technology
 
 The Marvin TUI is built with [Textual](https://github.com/Textualize/textual), a Python library for building TUIs.

mkdocs.yml

@@ -18,7 +18,7 @@ nav:
     - Concepts:
       - AI functions: guide/concepts/ai_functions.md
       - Bots: guide/concepts/bots.md
-      - ChatTUI: guide/concepts/tui.md
+      - TUI: guide/concepts/tui.md
       - Loaders: guide/concepts/loaders_and_documents.md
       - Infrastructure: guide/concepts/infra.md
       - Plugins: guide/concepts/plugins.md

src/marvin/bots/__init__.py

@@ -1,15 +1,21 @@
 import asyncio
 import marvin
 from . import fun, utilities, meta, jobs, assistants
+from marvin import get_logger
+
+logger = get_logger(__name__)
 
 
 async def install_bots():
     from marvin import Bot
 
-    for module in [fun, meta, jobs, assistants]:
-        for bot in module.__dict__.values():
-            if isinstance(bot, Bot):
-                await bot.save(if_exists="update")
+    try:
+        for module in [fun, meta, jobs, assistants]:
+            for bot in module.__dict__.values():
+                if isinstance(bot, Bot):
+                    await bot.save(if_exists="update")
+    except Exception:
+        logger.error("Failed to install bots", exc_info=True)
 
 
 if marvin.settings.bot_create_default_bots_on_startup:

src/marvin/cli/tui.py

@@ -496,10 +496,13 @@ def compose(self) -> ComposeResult:
         )
         yield Button("Upgrade DB", variant="warning", id="upgrade-db")
 
-    def on_button_pressed(self, event: Button.Pressed):
+    async def on_button_pressed(self, event: Button.Pressed):
         if event.button.id == "upgrade-db":
             event.button.label = "Upgrading..."
             marvin.infra.database.alembic_upgrade()
+            if marvin.settings.bot_create_default_bots_on_startup:
+                await marvin.bots.install_bots()
+            self.app.database_ready = True
             self.app.pop_screen()
 
 
@@ -734,6 +737,7 @@ class MarvinApp(App):
     )
     bot_responding: bool = reactive(False)
     is_ready: bool = False
+    database_ready: bool = reactive(False)
 
     def __init__(self, default_bot_name: str = None, **kwargs):
         super().__init__(**kwargs)
@@ -748,6 +752,7 @@ async def check_database_upgrade(self):
 
         try:
             await marvin.infra.database.check_alembic_version()
+            self.database_ready = True
         except marvin.infra.database.DatabaseWarning as w:
             if "Database migrations are not up to date" in str(w):
                 self.push_screen(DatabaseUpgradeScreen())
@@ -757,17 +762,20 @@ async def check_database_upgrade(self):
         # reset warning behavior
         warnings.resetwarnings()
 
+    async def watch_database_ready(self, database_ready: bool) -> None:
+        if database_ready:
+            try:
+                bot = await marvin.Bot.load(self._default_bot_name)
+            except HTTPException:
+                bot = marvin.bots.meta.marvin_bot
+            self.bot = bot
+
+            await self.bot.set_thread(self.thread_id)
+            self.is_ready = True
+
     async def on_ready(self) -> None:
         self.push_screen(MainScreen())
-        try:
-            bot = await marvin.Bot.load(self._default_bot_name)
-        except HTTPException:
-            bot = marvin.bots.meta.marvin_bot
-        self.bot = bot
-
-        await self.bot.set_thread(self.thread_id)
-        self.is_ready = True
-        self.set_timer(0.5, self.check_database_upgrade)
+        await self.check_database_upgrade()
 
     async def watch_bot(self, bot: marvin.Bot) -> None:
         if not self.is_ready: