SDK: use SQLite in-memory (off by default)

Author: stanislaw89

skyvern/cli/init_command.py

@@ -8,15 +8,15 @@
 from rich.progress import Progress, SpinnerColumn, TextColumn
 from rich.prompt import Confirm, Prompt
 
-from skyvern.forge.forge_app_initializer import start_forge_app
+from skyvern.forge.forge_app_initializer import setup_local_organization, start_forge_app
 from skyvern.utils import migrate_db
 from skyvern.utils.env_paths import resolve_backend_env_path
 
 from .browser import setup_browser_config
 from .console import console
 from .database import setup_postgresql
 from .llm_setup import setup_llm_providers, update_or_add_env_var
-from .mcp import setup_local_organization, setup_mcp
+from .mcp import setup_mcp
 
 
 def init_env(

skyvern/cli/mcp.py

@@ -6,49 +6,11 @@
 from rich.panel import Panel
 from rich.prompt import Confirm, Prompt
 
-from skyvern.forge import app
-from skyvern.forge.sdk.core import security
-from skyvern.forge.sdk.db.enums import OrganizationAuthTokenType
-from skyvern.forge.sdk.schemas.organizations import Organization
-from skyvern.forge.sdk.services.local_org_auth_token_service import SKYVERN_LOCAL_DOMAIN, SKYVERN_LOCAL_ORG
-from skyvern.forge.sdk.services.org_auth_token_service import API_KEY_LIFETIME
 from skyvern.utils import detect_os, get_windows_appdata_roaming
 from skyvern.utils.env_paths import resolve_backend_env_path
 
 from .console import console
 
-
-async def get_or_create_local_organization() -> Organization:
-    organization = await app.DATABASE.get_organization_by_domain(SKYVERN_LOCAL_DOMAIN)
-    if not organization:
-        organization = await app.DATABASE.create_organization(
-            organization_name=SKYVERN_LOCAL_ORG,
-            domain=SKYVERN_LOCAL_DOMAIN,
-            max_steps_per_run=10,
-            max_retries_per_step=3,
-        )
-        api_key = security.create_access_token(
-            organization.organization_id,
-            expires_delta=API_KEY_LIFETIME,
-        )
-        # generate OrganizationAutoToken
-        await app.DATABASE.create_org_auth_token(
-            organization_id=organization.organization_id,
-            token=api_key,
-            token_type=OrganizationAuthTokenType.api,
-        )
-    return organization
-
-
-async def setup_local_organization() -> str:
-    organization = await get_or_create_local_organization()
-    org_auth_token = await app.DATABASE.get_valid_org_auth_token(
-        organization_id=organization.organization_id,
-        token_type=OrganizationAuthTokenType.api.value,
-    )
-    return org_auth_token.token if org_auth_token else ""
-
-
 # ----- Helper paths and checks -----
 
 

skyvern/forge/api_app.py

@@ -16,6 +16,7 @@
 from skyvern.config import settings
 from skyvern.exceptions import SkyvernHTTPException
 from skyvern.forge import app as forge_app
+from skyvern.forge.forge_app import ForgeApp
 from skyvern.forge.forge_app_initializer import start_forge_app
 from skyvern.forge.request_logging import log_raw_request_middleware
 from skyvern.forge.sdk.core import skyvern_context
@@ -79,7 +80,10 @@ def create_api_app() -> FastAPI:
     """
 
     forge_app_instance = start_forge_app()
+    return create_fast_api_app(forge_app_instance)
 
+
+def create_fast_api_app(forge_app_instance: ForgeApp) -> FastAPI:
     fastapi_app = FastAPI(lifespan=lifespan)
 
     # Add CORS middleware

skyvern/forge/forge_app.py

@@ -84,15 +84,15 @@ class ForgeApp:
     agent: ForgeAgent
 
 
-def create_forge_app() -> ForgeApp:
+def create_forge_app(db: AgentDB | None = None) -> ForgeApp:
     """Create and initialize a ForgeApp instance with all services"""
     settings: Settings = SettingsManager.get_settings()
 
     app = ForgeApp()
 
     app.SETTINGS_MANAGER = settings
 
-    app.DATABASE = AgentDB(settings.DATABASE_STRING, debug_enabled=settings.DEBUG_MODE)
+    app.DATABASE = db if db else AgentDB(settings.DATABASE_STRING, debug_enabled=settings.DEBUG_MODE)
     if settings.SKYVERN_STORAGE_TYPE == "s3":
         StorageFactory.set_storage(S3Storage())
     app.STORAGE = StorageFactory.get_storage()

skyvern/forge/forge_app_initializer.py

@@ -1,14 +1,20 @@
 import structlog
 
 from skyvern.config import settings
-from skyvern.forge import set_force_app_instance
+from skyvern.forge import app, set_force_app_instance
 from skyvern.forge.forge_app import ForgeApp, create_forge_app
+from skyvern.forge.sdk.core import security
+from skyvern.forge.sdk.db.client import AgentDB
+from skyvern.forge.sdk.db.enums import OrganizationAuthTokenType
+from skyvern.forge.sdk.schemas.organizations import Organization
+from skyvern.forge.sdk.services.local_org_auth_token_service import SKYVERN_LOCAL_DOMAIN, SKYVERN_LOCAL_ORG
+from skyvern.forge.sdk.services.org_auth_token_service import API_KEY_LIFETIME
 
 LOG = structlog.get_logger()
 
 
-def start_forge_app() -> ForgeApp:
-    force_app_instance = create_forge_app()
+def start_forge_app(db: AgentDB | None = None) -> ForgeApp:
+    force_app_instance = create_forge_app(db)
     set_force_app_instance(force_app_instance)
 
     if settings.ADDITIONAL_MODULES:
@@ -26,3 +32,34 @@ def start_forge_app() -> ForgeApp:
         )
 
     return force_app_instance
+
+
+async def get_or_create_local_organization() -> Organization:
+    organization = await app.DATABASE.get_organization_by_domain(SKYVERN_LOCAL_DOMAIN)
+    if not organization:
+        organization = await app.DATABASE.create_organization(
+            organization_name=SKYVERN_LOCAL_ORG,
+            domain=SKYVERN_LOCAL_DOMAIN,
+            max_steps_per_run=10,
+            max_retries_per_step=3,
+        )
+        api_key = security.create_access_token(
+            organization.organization_id,
+            expires_delta=API_KEY_LIFETIME,
+        )
+        # generate OrganizationAutoToken
+        await app.DATABASE.create_org_auth_token(
+            organization_id=organization.organization_id,
+            token=api_key,
+            token_type=OrganizationAuthTokenType.api,
+        )
+    return organization
+
+
+async def setup_local_organization() -> str:
+    organization = await get_or_create_local_organization()
+    org_auth_token = await app.DATABASE.get_valid_org_auth_token(
+        organization_id=organization.organization_id,
+        token_type=OrganizationAuthTokenType.api.value,
+    )
+    return org_auth_token.token if org_auth_token else ""

skyvern/forge/sdk/db/client.py

@@ -3,7 +3,7 @@
 from typing import Any, List, Literal, Sequence, overload
 
 import structlog
-from sqlalchemy import and_, asc, case, delete, distinct, exists, func, or_, pool, select, tuple_, update
+from sqlalchemy import and_, asc, case, delete, exists, func, or_, pool, select, update
 from sqlalchemy.exc import SQLAlchemyError
 from sqlalchemy.ext.asyncio import AsyncEngine, async_sessionmaker, create_async_engine
 

skyvern/library/embedded_server_factory.py

@@ -2,17 +2,24 @@
 from typing import Any
 
 import httpx
+import structlog
+from dotenv import load_dotenv
 from httpx import ASGITransport
 
 from skyvern.config import settings
-from skyvern.forge.api_app import create_api_app
+from skyvern.forge.api_app import create_api_app, create_fast_api_app
+from skyvern.forge.forge_app_initializer import setup_local_organization, start_forge_app
 from skyvern.forge.sdk.api.llm.config_registry import LLMConfigRegistry
 from skyvern.forge.sdk.api.llm.models import LLMConfig, LLMRouterConfig
+from skyvern.forge.sdk.db.client import AgentDB
+
+LOG = structlog.get_logger()
 
 
 def create_embedded_server(
     llm_config: LLMRouterConfig | LLMConfig | None = None,
     settings_overrides: dict[str, Any] | None = None,
+    use_in_memory_db: bool = False,
 ) -> httpx.AsyncClient:
     class EmbeddedServerTransport(httpx.AsyncBaseTransport):
         def __init__(self) -> None:
@@ -21,6 +28,8 @@ def __init__(self) -> None:
 
         async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
             if self._transport is None:
+                load_dotenv(".env")
+
                 settings.BROWSER_LOGS_ENABLED = False
 
                 if llm_config:
@@ -38,10 +47,28 @@ async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
                         else:
                             raise ValueError(f"Invalid setting: {key}")
 
-                self._api_key = os.getenv("SKYVERN_API_KEY")
-                if not self._api_key:
-                    raise ValueError("SKYVERN_API_KEY is not set. Provide api_key or set SKYVERN_API_KEY in .env file.")
-                api_app = create_api_app()
+                if use_in_memory_db:
+                    from sqlalchemy.ext.asyncio import create_async_engine  # noqa: PLC0415
+
+                    from skyvern.forge.sdk.db.models import Base  # noqa: PLC0415
+
+                    settings.DATABASE_STRING = "sqlite+aiosqlite:///:memory:"
+                    engine = create_async_engine(settings.DATABASE_STRING)
+                    async with engine.begin() as conn:
+                        await conn.run_sync(Base.metadata.create_all)
+
+                    db = AgentDB(settings.DATABASE_STRING, debug_enabled=settings.DEBUG_MODE, db_engine=engine)
+                    forge_app_instance = start_forge_app(db)
+                    api_app = create_fast_api_app(forge_app_instance)
+                    self._api_key = await setup_local_organization()
+                    LOG.info("Embedded server initialized with in-memory database")
+                else:
+                    self._api_key = os.getenv("SKYVERN_API_KEY")
+                    if not self._api_key:
+                        raise ValueError(
+                            "SKYVERN_API_KEY is not set. Provide api_key or set SKYVERN_API_KEY in .env file."
+                        )
+                    api_app = create_api_app()
 
                 self._transport = ASGITransport(app=api_app)
 

skyvern/library/skyvern.py

@@ -128,6 +128,7 @@ def local(
         *,
         llm_config: LLMRouterConfig | LLMConfig | None = None,
         settings: dict[str, Any] | None = None,
+        use_in_memory_db: bool | None = None,
     ) -> "Skyvern":
         """Local/embedded mode: Run Skyvern locally in-process.
 
@@ -186,19 +187,32 @@ def local(
             settings: Optional dictionary of Skyvern settings to override.
                 These override the corresponding settings from your .env file.
                 Example: {"MAX_STEPS_PER_RUN": 100, "BROWSER_TYPE": "chromium-headful"}
+            use_in_memory_db: Optional flag to use an in-memory database.
 
         Returns:
             Skyvern: A Skyvern instance running in local/embedded mode.
         """
         from skyvern.library.embedded_server_factory import create_embedded_server  # noqa: PLC0415
 
-        if not os.path.exists(".env"):
-            raise ValueError("Please run `skyvern quickstart` to set up your local Skyvern environment")
-
         load_dotenv(".env")
-        api_key = os.getenv("SKYVERN_API_KEY")
-        if not api_key:
-            raise ValueError("SKYVERN_API_KEY is not set. Provide api_key or set SKYVERN_API_KEY in .env file.")
+
+        if use_in_memory_db is not None:
+            do_use_in_memory_db = use_in_memory_db
+        else:
+            do_use_in_memory_db = os.environ.get("SKYVERN_USE_IN_MEMORY_DB", "").lower() == "true"
+
+        # Validate prerequisites based on configuration mode
+        if not do_use_in_memory_db and not os.path.exists(".env"):
+            raise ValueError("Please run `skyvern quickstart` to set up your local Skyvern environment.")
+
+        if not llm_config and not os.path.exists(".env"):
+            raise ValueError(
+                "LLM configuration is required. Either provide llm_config parameter or "
+                "run `skyvern quickstart` to configure LLM settings in .env file."
+            )
+
+        if not do_use_in_memory_db and not os.getenv("SKYVERN_API_KEY"):
+            raise ValueError("Please run `skyvern quickstart` to set up your local Skyvern environment.")
 
         obj = cls.__new__(cls)
 
@@ -208,6 +222,7 @@ def local(
             httpx_client=create_embedded_server(
                 llm_config=llm_config,
                 settings_overrides=settings,
+                use_in_memory_db=do_use_in_memory_db,
             ),
         )